pose encoder trial

scripts/pose_encoder_train.py

@@ -1,14 +1,10 @@
 import os
 import sys
-
-os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
-os.environ["CUDA_VISIBLE_DEVICES"] = '4' #config['gpu']
-import torch
-
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 
-import importlib
 import argparse
+import importlib
+
 parser = argparse.ArgumentParser(description='')
 parser.add_argument('--name', type=str)
 parser.add_argument('--config', type=str)
@@ -17,9 +13,16 @@ args = vars(parser.parse_args())
 
 # config_module = importlib.import_module(args['config'])
 # config = config_module.config
+if args['debug']:
+    pass
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+if args['debug']:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 
 import wandb
+import torch
 import torch.optim as optim
 import torch.nn as nn
 from torch.utils.data import DataLoader
@@ -37,15 +40,11 @@ ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
 
 
 def init_wandb():
-    run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
+    run = wandb.init(project='reconstruction', name='pose_1', notes='', dir='/data-local/127/temp', tags=['pose'])
     artifact = wandb.Artifact('code', type='code')
-    artifact.add_file('scripts/cvae.py')
-    artifact.add_file('losses/losses.py')
-    artifact.add_file('models/base.py')
-    artifact.add_file('models/cond_encoder.py')
-    artifact.add_file('models/convnext.py')
-    artifact.add_file('models/cvae_model.py')
-    artifact.add_file('data/dataset.py')
+    artifact.add_file('scripts/pose_encoder_train.py')
+    artifact.add_file('models/pose_encoder.py')
+    artifact.add_file('data/pose_dataset.py')
     config_file_names = [key for key in sys.modules.keys()  if key.startswith('configs.')]
     for config_file_name in config_file_names:
         file_name = '/'.join(config_file_name.split('.')) + '.py'