ae_curv for cape

configs/ae_curv_cape_m.py

@@ -0,0 +1,72 @@
+from utils.convnext_utils import CNBlockConfig
+
+config = dict()
+
+config['gpu'] = '5'
+config['name'] = 'curv_torch_cape_m'
+config['notes'] = ' ' 
+config['tags'] = ['torch', 'curv'] 
+config['script'] = 'scripts/ae_curv_train.py'
+
+config['pretrained_model'] = ''
+
+config['batch_size'] = 4
+config['train_dir'] = '/data-local/127/cape_processed/train/male'
+config['valid_dir'] = '/data-local/127/cape_processed/test/male'
+config['test_dir'] = '/data-local/127/cape_processed/test/male'
+config['train_frames'] = None
+config['valid_frames'] = None
+
+config['resolution'] = 512
+config['is_skirt'] = False
+config['dataset_in_memory'] = False
+config['is_cape'] = True
+config['remove_global_orientation'] = True
+config['device'] = 'cuda'
+
+config['smpl_res'] = 14475
+if config['is_cape']:
+  config['smpl_res'] = 3638
+
+config['latent_dim'] = 128
+
+config['bottleneck_scale'] = 2
+
+config['cond_ae'] = dict()
+config['cond_ae']['device'] = config['device']
+config['cond_ae']['is_label'] = True
+config['cond_ae']['latent_dim'] = config['latent_dim']
+config['cond_ae']['kernel_size'] = 7
+config['cond_ae']['bottleneck_size'] = 40 * config['bottleneck_scale']
+
+block_setting = [
+    CNBlockConfig(16, 32, 3),
+    CNBlockConfig(32, 64, 3),
+    CNBlockConfig(64, 128, 9),
+    CNBlockConfig(128, config['cond_ae']['bottleneck_size'], 9),
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3),
+]
+
+block_setting_reverse = [
+    CNBlockConfig(config['cond_ae']['bottleneck_size'], 128, 3),
+    CNBlockConfig(128, 64, 9),
+    CNBlockConfig(64, 32, 9),
+    CNBlockConfig(32, 16, 3),
+    CNBlockConfig(16, 3, 3),
+]
+
+config['cond_ae']['encoder_block_settings'] = block_setting
+config['cond_ae']['decoder_block_settings'] = block_setting_reverse
+config['cond_ae']['activation'] = 'tanh'
+
+
+config['lr'] = 0.001
+config['cond_ae']['upscale_output'] = False
+
+config['signature'] = { 'id',
+                        'frame_nr',
+                        'uv_curvature',
+                        'mesh_curvature',
+                        'mesh_mask',
+                        'uv_mask',
+                        }

scripts/ae_curv_train.py

@@ -8,14 +8,17 @@ import importlib
 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
-# from configs.ae_curv_1 import config
+if args['debug']:
+    from configs.ae_curv_cape_m 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 wandb
@@ -160,37 +163,31 @@ def train():
             if step == 0:
                 to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
 
+            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_enc(model.encoder, step, optimizer, None)
-            ckpt_saver_dec(model.decoder, step, optimizer)
+            ckpt_saver_dec(model.decoder, step, optimizer, None)
 
         step += 1
 
 
 if __name__ == '__main__':
+    if not args['debug']:
         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())
-    
-    resolution_modifier = ''
-    if config['resolution'] != 512:
-        resolution_modifier = '_' + str(config['resolution'])
     skirt_modifier = '_skirt' if config['is_skirt'] else ''
 
     uv_mask_body = np.load('custom_utils/mask_256x512' + ('_skirt' if config['is_skirt'] else '') + '.npy')
     uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
     uv_mask_body = torch.from_numpy(uv_mask_body).float().to(config['device'])
 
-    dataset_train = Dataset(config['train_dir'], config['signature'], device=config['device'], frames_file=config['train_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'])
+    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()
@@ -210,11 +207,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_enc = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/' + config['name'] + '_dec', 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()