update training scripts

scripts/ae_static_train.py

@@ -13,8 +13,8 @@ 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
+    # from configs.ae_static_cape_m import config
+    from configs.iccvw.ae_static_1 import config
 os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
 if args['debug']:
@@ -49,6 +49,7 @@ from pytorch_msssim import MS_SSIM
 
 
 def init_wandb():
+    wandb.login(host='https://wandb2.lnrweb.com', key='local-2a78b434e9bc51e9f12625f8321a9e3ee34af16d', relogin=True, force=True)
     run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
     artifact = wandb.Artifact('code', type='code')
     artifact.add_file('scripts/ae_static_train.py')
@@ -208,15 +209,15 @@ 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/torch/' + config['name'] + '_enc', top_n=5)
-    ckpt_saver_dec = CheckpointSaver('checkpoints/iccv/torch/' + config['name'] + '_dec', top_n=5)
+    ckpt_saver_enc = CheckpointSaver('checkpoints/iccvw/' + config['name'] + '_enc', top_n=5)
+    ckpt_saver_dec = CheckpointSaver('checkpoints/iccvw/' + 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)
+        wandb.watch(model, log='all', log_freq=10000, log_graph=True)
 
     train()