diff --git a/configs/ae_static_cape_m.py b/configs/ae_static_cape_m.py new file mode 100644 index 0000000..12499b1 --- /dev/null +++ b/configs/ae_static_cape_m.py @@ -0,0 +1,74 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '5' +config['name'] = 'static_torch_cape_m' +config['notes'] = ' ' +config['tags'] = ['torch', 'static'] +config['script'] = 'scripts/ae_static_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'] = True +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'] = False +config['cond_ae']['latent_dim'] = config['latent_dim'] +config['cond_ae']['kernel_size'] = 7 +config['cond_ae']['bottleneck_size'] = 8 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(8, 16, 3), + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, config['cond_ae']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), + # CNBlockConfig(3, 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.0005 +config['cond_ae']['upscale_output'] = False +config['pretrained_model'] = '' + +config['signature'] = { 'id', + 'frame_nr', + 'uv_static', + 'mesh_static', + 'mesh_mask', + 'uv_mask', + } \ No newline at end of file diff --git a/scripts/ae_static_train.py b/scripts/ae_static_train.py index ee693f8..1d03715 100644 --- a/scripts/ae_static_train.py +++ b/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()