fix upscaler

scripts/upscaler_train.py

@@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '7'
+os.environ["CUDA_VISIBLE_DEVICES"] = '6'
 
 
 import time
@@ -26,6 +26,7 @@ import plotly.graph_objects as go
 import plotly.express as ex
 import torch.optim as optim
 import torch.nn.functional as F
+import torchvision
 
 from data.dataset import Dataset
 # from torch.utils.data import DataLoader
@@ -121,22 +122,26 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
 
 def train_step(batch):
     model.train()
-    with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
-        pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
-        
-        l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
-        if low_res == 32:
-            ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
-        else:
-            ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
+    pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
+    # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+    pred = resizer(pred)
+    
+    # l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+    # loss = l1_loss
+    if low_res == 32:
+        ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    else:
+        ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
+    loss = 1e3 * (1 - ms_ssim)
 
-        loss = 1e3 * (1 - ms_ssim)
-
-    scaler.scale(loss).backward()
-    scaler.step(optimizer)
-    scaler.update()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
     optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
 
     model.eval()
     with torch.inference_mode():
@@ -144,7 +149,7 @@ def train_step(batch):
         mesh_orig = batch['mesh_unposed']
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
 
-    return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, }
+    return { 'train/loss': loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } # 'train/l1_loss': l1_loss,
         
 
 def valid_step(batch):
@@ -152,27 +157,30 @@ def valid_step(batch):
 
     with torch.inference_mode():
         pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+        # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+        pred = resizer(pred)
         
-        l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+        # l1_loss = torch.mean(torch.sum(torch.abs(batch['uv_unposed'] - pred * batch['uv_mask'].unsqueeze(1)), dim=-1))
+        # loss = l1_loss
         if low_res == 32:
             ms_ssim = torch.mean(_MSSSIM_small(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
         else:
             ms_ssim = torch.mean(_MSSSIM(pred * batch['uv_mask'].unsqueeze(1) + 1, batch['uv_unposed'] * batch['uv_mask'].unsqueeze(1) + 1))
-
         loss = 1e3 * (1 - ms_ssim)
+
         mesh_pred = uv2mesh(pred)
         mesh_orig = batch['mesh_unposed']
         metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'].unsqueeze(2) * (mesh_orig - mesh_pred) ** 2, dim=-1))) / torch.sum(batch['mesh_mask'])
 
-    return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, }
+    return { 'valid/loss': loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } #  'valid/l1_loss': l1_loss,
        
 
 def plot_everything(batch: dict, type: str):
     model.eval()
     with torch.inference_mode():
         pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]])
-        pred = F.interpolate(pred, size=(512, 256), mode='bilinear')
+        # pred = F.interpolate(pred, size=(512, 256), mode='bilinear', align_corners=True, antialias=True)
+        pred = resizer(pred)
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_unposed']
@@ -187,7 +195,7 @@ def plot_everything(batch: dict, type: str):
     # figs = { type + "/garment_" + str(i): plot_3d(gt[i], garment_unposed[i], faces[i], collision[i], collision_faces[i], collision_color[i]) for i in range(4) }
     
     orig_img = batch['uv_unposed']
-    img = torch.concat((orig_img, pred), axis=2)
+    img = torch.concat((orig_img, pred), axis=3)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
 
     return { **figs, **{ type + '/images': imgs} }
@@ -221,8 +229,8 @@ def train():
             valid_accumulator.reset()
             print(step)
 
-        if step % 10000 == 0:
-            ckpt_saver(model, step, optimizer, scaler)
+        if step % 10000 == 0 and not args['debug']:
+            ckpt_saver(model, step, optimizer, None)
 
         step += 1
 
@@ -263,10 +271,11 @@ if __name__ == '__main__':
         
     model : Upscaler = Upscaler(kernel_size=kernel_size).to('cuda')
     optimizer = optim.Adam(model.parameters(), lr=0.0001)
-    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    # scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    resizer = torchvision.transforms.Resize((512, 256), interpolation=torchvision.transforms.InterpolationMode.BILINEAR)
 
     # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
-    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.5, 0.2, 0.2, 0.1])
+    _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[0.8, 0.2])
     _MSSSIM_small = MS_SSIM(data_range=2, size_average=True, channel=3, weights=[1.0])
 
     step = 0