ae_static for cape

scripts/ae_static_train.py

@@ -7,14 +7,18 @@ import argparse
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
+parser.add_argument('--debug', action='store_true')
 args = vars(parser.parse_args())
 
 config_module = importlib.import_module(args['config'])
 config = config_module.config
-
+if args['debug']:
+    from configs.ae_static_cape_m import config
+    # from configs.ae_static_orig_2 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
-
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 import time
 import argparse
@@ -66,6 +70,7 @@ uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_
 
 def train_step(batch):
     model.train()
+    # with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=use_amp):
     pred = model([batch['uv_static'], batch['uv_mask']])
 
     ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1))
@@ -75,6 +80,10 @@ def train_step(batch):
     optimizer.zero_grad()
     loss.backward()
     optimizer.step()
+    # scaler.scale(loss).backward()
+    # scaler.step(optimizer)
+    # scaler.update()
+    # optimizer.zero_grad()
 
     model.eval()
     with torch.inference_mode():
@@ -113,7 +122,10 @@ def plot_everything(batch: dict, type: str):
     
     mesh_pred = uv2mesh(pred)
     mesh_orig = batch['mesh_static']
-    
+
+    mesh_pred = torch.cat((mesh_pred[:, :, 0:1], -mesh_pred[:, :, 2:3], mesh_pred[:, :, 1:2]), dim=2)
+    mesh_orig = torch.cat((mesh_orig[:, :, 0:1], -mesh_orig[:, :, 2:3], mesh_orig[:, :, 1:2]), dim=2)
+
     orig_img = torch.permute(batch['uv_static'], [0,2,3,1])
     pred_img = torch.permute(pred, [0,2,3,1])
     img = torch.concat((orig_img, pred_img), axis=2)
@@ -157,34 +169,29 @@ def train():
             if step == 0:
                 to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
-            wandb.log(to_log, step=step)
+            if not args['debug']:
+                wandb.log(to_log, step=step)
             train_accumulator.reset()
             valid_accumulator.reset()
             print(step)
 
         if step % 10000 == 0:
-            ckpt_saver_enc(model.encoder, step, optimizer)
-            ckpt_saver_dec(model.decoder, step, optimizer)
+            ckpt_saver_enc(model.encoder, step, optimizer, scaler)
+            ckpt_saver_dec(model.decoder, step, optimizer, scaler)
 
         step += 1
 
 
 if __name__ == '__main__':
-    init_wandb()
-    # device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # parser = argparse.ArgumentParser(description='')
-    # parser.add_argument('--name', type=str)
-    # args = vars(parser.parse_args())
+    if not args['debug']:
+        init_wandb()
     
-    resolution_modifier = ''
-    if config['resolution'] != 512:
-        resolution_modifier = '_' + str(config['resolution'])
+    use_amp = True
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     # dataset = Dataset('/data-local/127/train', config['signature'], resolutions=[512], device=device, frames_file='frames_jumpsuit.txt')
-    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'])
-    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'])
+    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
+    dataset_valid = Dataset(config['valid_dir'], config['signature'], device=config['device'], frames_file=config['valid_frames'], use_cache=config['dataset_in_memory'], is_cape=config['is_cape'])
     train_dataloader = DataLoader(dataset_train, batch_size=4, drop_last=True, shuffle=True, )
     valid_dataloader = DataLoader(dataset_valid, batch_size=4, drop_last=True, shuffle=True)
     train_accumulator = Accumulator()
@@ -192,6 +199,7 @@ if __name__ == '__main__':
 
     model : AutoEncoder = AutoEncoder(config['cond_ae'])
     optimizer = optim.Adam(model.parameters(), lr=config['lr'])
+    scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
 
     # _MSSSIM = MSSSIM(data_range=2).to(config['device'])
     _MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
@@ -200,11 +208,16 @@ if __name__ == '__main__':
         load_model(model.encoder, config['pretrained_model'] + '_enc')
         load_model(model.decoder, config['pretrained_model'] + '_dec')
 
-    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', top_n=5)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
 
     run_name = config['name']
     print(f'started training: {run_name}')
+
+    if not args['debug']:
+        # model = torch.compile(model)
+        wandb.watch(model, log='all', log_freq=1000, log_graph=True)
+
     train()