luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

extrapolation train

Browse Source
This commit is contained in:
luni
2023-05-25 18:09:55 +02:00
parent 843c9ac4b0
commit af9811e731
7 changed files with 596 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/cvae.py
+10 -9
View File
@@ -10,10 +10,12 @@ 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']:
pass
from configs.cape_00032_SL import config
else:
config_module = importlib.import_module(args['config'])
config = config_module.config
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
if args['debug']:
@@ -50,7 +52,7 @@ from custom_utils.custom_utils import UV2Mesh, Accumulator, plot_3d, get_summary
def init_wandb():
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/data-local/127/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('scripts/cvae.py')
artifact.add_file('losses/losses.py')
@@ -457,7 +459,7 @@ def train():
g_balancing_coeff += 1 / 10000
if step % 10000 == 0:
ckpt_saver(model, step, optimizer)
ckpt_saver(model, step, optimizer, None)
step += 1
@@ -469,8 +471,8 @@ if __name__ == '__main__':
smpl = SMPL(orig_smpl=config['is_cape'], path='smpl', device=config['device'])
skirt_modifier = '_skirt' if config['is_skirt'] else ''
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()
@@ -478,12 +480,11 @@ if __name__ == '__main__':
model = CVAE(config)
optimizer = optim.Adam(model.parameters(), lr=config['lr'])
_MSSSIM = MS_SSIM(data_range=2, size_average=True, channel=3)
g_balancing_coeff = config['starting_balancing_coeff']
ckpt_saver = CheckpointSaver('checkpoints/iccv/' + config['name'], top_n=5)
ckpt_saver = CheckpointSaver('checkpoints/iccv/torch/' + config['name'], top_n=10)
step = 0
if config['pretrained_model'] != '':