update

scripts/cvae_pyramid_fine.py

@@ -40,6 +40,8 @@ import torch
 import torch.nn as nn
 import torch.optim as optim
 import torch.nn.functional as F
+from torchvision.transforms import Resize
+from torchvision.transforms.functional import InterpolationMode
 from data.dataset import Dataset
 from utils.checkpoints import CheckpointSaver, load_model
 from utils.utils import InfiniteDataLoader as DataLoader
@@ -57,7 +59,7 @@ def init_wandb():
     wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     run = wandb.init(project='iccvw', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('scripts/cvae_pyramid.py')
+    artifact.add_file('scripts/cvae_pyramid_fine.py')
     artifact.add_file('losses/losses.py')
     artifact.add_file('models/base.py')
     artifact.add_file('models/cond_encoder.py')
@@ -241,8 +243,10 @@ def train_step(batch):
         res = config['pyramid']['resolutions'][i]
         res_modifier = '_' + str(res) if res != 512 else ''
         pred, mean, logvar = out_levels[i]
+        if res != 512:
+            pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
-        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((levels_input[i][0] - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
+        l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) # levels_input[i][0] = uv_orig
         kl = kl_loss(mean, logvar)
 
         mesh_pred = uv2mesh(pred)
@@ -281,7 +285,7 @@ def train_step(batch):
         mean = torch.mean(mean)
         logvar = torch.mean(logvar)
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, levels_input[i][0] + 1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
     output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 
@@ -318,6 +322,8 @@ def valid_step(batch):
             res = config['pyramid']['resolutions'][i]
             res_modifier = '_' + str(res) if res != 512 else ''
             pred, mean, logvar = out_levels[i]
+            if res != 512:
+                pred = resizer(pred) * batch['uv_mask'][:,None,:,:]
 
             mesh_pred = uv2mesh(pred)
 
@@ -336,7 +342,7 @@ def valid_step(batch):
             c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio)
             l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1))
             n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_fix']['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
-            l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((levels_input[i][0] - pred) ** 2), dim=-1), dim=[1, 2]))
+            l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2]))
             
             kl = kl_loss(mean, logvar)
 
@@ -355,7 +361,7 @@ def valid_step(batch):
         mean = torch.mean(mean)
         logvar = torch.mean(logvar)
 
-        ms_ssim = torch.mean(_MSSSIM(pred + 1, levels_input[i][0] + 1))
+        ms_ssim = torch.mean(_MSSSIM(pred + 1, uv_orig + 1))
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_fix']['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / torch.sum(batch['mesh_fix']['mesh_mask'])
 
         output = { **output, **{ 'valid/edge_loss' + res_modifier: e_loss, 'valid/collision_loss' + res_modifier: c_loss,'valid/euclid_error' + res_modifier: metric, 
@@ -538,15 +544,18 @@ if __name__ == '__main__':
 
     model: Pyramid = Pyramid(config)
 
-    optimizer = optim.Adam(model.parameters(), lr=config['lr'] * 0.05)
+    # optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    optimizer = optim.Adam(model.parameters(), lr=config['lr'] * 0.1)
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
-    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_fine_2', top_n=50)
+    ckpt_saver = CheckpointSaver('checkpoints/iccvw/' + config['name'], top_n=50)
     scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    resizer = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device'])
 
     g_balancing_coeff = config['starting_balancing_coeff']
-    config['pyramid']['loss_weights'] = [1,3,5,10,25]
+    config['pyramid']['loss_weights'] = [1,2,3,4,5]
 
-    step = load_model(model, f'checkpoints/iccvw/{config["name"]}', iteration='25000')
+    step = 0
+    load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration'])
 
     for param in model.static_enc.parameters(): param.requires_grad = False
     for param in model.body_enc.parameters(): param.requires_grad = False