diff --git a/configs/abl_all_2.py b/configs/abl_all_2.py new file mode 100644 index 0000000..e4458b5 --- /dev/null +++ b/configs/abl_all_2.py @@ -0,0 +1,192 @@ +from utils.convnext_utils import CNBlockConfig + +config = dict() + +config['gpu'] = '6' +config['name'] = 'abl_torch_all_2' +config['notes'] = '' +config['tags'] = ['iccv', 'abl_all', 'torch'] +config['script'] = 'scripts/cvae.py' + +config['pretrained_model'] = 'checkpoints/iccv/abl_torch_all' + +config['use_skinning'] = True # and deshaping +config['remove_global_orientation'] = True +config['use_curvature'] = True +config['use_extra_losses'] = True # edge, collision, 3d mesh +config['use_offsets'] = True + +config['use_norm_loss'] = True + +ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False) +assert ablation_count == 1 or ablation_count == 0 + +config['avg_garment_path'] = 'custom_data/avg/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv train_no_shape_one_hot_avg_uv +if config['use_skinning']: + assert (config['use_offsets'] and config['avg_garment_path'] != '') or (not config['use_offsets'] and config['avg_garment_path'] == '') + +config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision + +config['batch_size'] = 4 +config['resolution'] = 512 +config['is_skirt'] = False +config['dataset_in_memory'] = False +config['is_cape'] = False +config['device'] = 'cuda' + +config['train_dir'] = '/data-local/127/train' +config['valid_dir'] = '/data-local/127/valid' +config['test_dir'] = '/data-net/dataset/test_ready' +config['train_frames'] = None +config['valid_frames'] = None + +config['static_model'] = 'checkpoints/iccv/static_torch_orig_2' +config['body_model'] = 'checkpoints/iccv/body_torch_1' +config['curvature_model'] = 'checkpoints/iccv/curv_torch_1' + +config['smpl_res'] = 14475 +if config['is_cape']: + config['smpl_res'] = 3638 + +config['latent_dim'] = 128 + +############### +config['lr'] = 0.0001# 0.0001 +############### +config['starting_balancing_coeff'] = 0.0001 + +config['signature']: set = { 'id', + 'frame_nr', + 'pose_shape', + 'uv_body_posed', + 'mesh_body_posed', + 'mesh_posed', + 'uv_curvature', + 'uv_static', + 'mesh_mask', + 'uv_mask', + } + +if config['use_skinning']: + config['signature'].add('uv_unposed') + config['signature'].add('mesh_unposed') +else: + config['signature'].add('uv_posed') + +config['upscale_enhancer'] = '' +config['freeze_encoders'] = True +config['use_one_hot'] = False +config['bottleneck_scale'] = 2 + +config['cvae'] = dict() +config['cvae']['device'] = config['device'] +config['cvae']['is_label'] = False +config['cvae']['input_features'] = 3+3+3+7 if config['use_curvature'] else 3+3+3 +config['cvae']['kernel_size'] = 7 +config['cvae']['bottleneck_size'] = 64 * config['bottleneck_scale'] +config['cvae']['latent_dim'] = config['cvae']['bottleneck_size'] + +block_setting = [ + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 3), + CNBlockConfig(128, 256, 3), + CNBlockConfig(256, config['cvae']['bottleneck_size'], 9), + CNBlockConfig(config['cvae']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cvae']['bottleneck_size'] * 2, 256, 3), + CNBlockConfig(256, 128, 9), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 3), + CNBlockConfig(32, 3, 3), +] + +config['cvae']['encoder_block_settings'] = block_setting +config['cvae']['decoder_block_settings'] = block_setting_reverse + + + +####################### cond_ae_curv ############################ +config['cond_ae_curv'] = dict() +config['cond_ae_curv']['device'] = config['device'] +config['cond_ae_curv']['is_label'] = True +config['cond_ae_curv']['latent_dim'] = config['latent_dim'] +config['cond_ae_curv']['kernel_size'] = 7 +config['cond_ae_curv']['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_curv']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_curv']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_curv']['encoder_block_settings'] = block_setting +config['cond_ae_curv']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_curv']['activation'] = 'tanh' + +####################### cond_ae_body ############################ +config['cond_ae_body'] = dict() +config['cond_ae_body']['device'] = config['device'] +config['cond_ae_body']['is_label'] = False +config['cond_ae_body']['latent_dim'] = config['latent_dim'] +config['cond_ae_body']['kernel_size'] = 7 +config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale'] + +block_setting = [ + CNBlockConfig(16, 32, 3), + CNBlockConfig(32, 64, 3), + CNBlockConfig(64, 128, 9), + CNBlockConfig(128, config['cond_ae_body']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_body']['bottleneck_size'], 128, 3), + CNBlockConfig(128, 64, 9), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 3, 3), +] + +config['cond_ae_body']['encoder_block_settings'] = block_setting +config['cond_ae_body']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_body']['activation'] = 'tanh' + +####################### cond_ae_static ############################ +config['cond_ae_static'] = dict() +config['cond_ae_static']['device'] = config['device'] +config['cond_ae_static']['is_label'] = False +config['cond_ae_static']['latent_dim'] = config['latent_dim'] +config['cond_ae_static']['kernel_size'] = 7 +config['cond_ae_static']['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_static']['bottleneck_size'], 9), + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], None, 3), +] + +block_setting_reverse = [ + CNBlockConfig(config['cond_ae_static']['bottleneck_size'], 64, 3), + CNBlockConfig(64, 32, 9), + CNBlockConfig(32, 16, 3), + CNBlockConfig(16, 8, 3), + CNBlockConfig(8, 3, 3), +] + +config['cond_ae_static']['encoder_block_settings'] = block_setting +config['cond_ae_static']['decoder_block_settings'] = block_setting_reverse +config['cond_ae_static']['activation'] = 'tanh' \ No newline at end of file diff --git a/scripts/ae_body_train.py b/scripts/ae_body_train.py index 0cf2bff..321f4b1 100644 --- a/scripts/ae_body_train.py +++ b/scripts/ae_body_train.py @@ -13,7 +13,8 @@ args = vars(parser.parse_args()) config_module = importlib.import_module(args['config']) config = config_module.config -# from configs.ae_body_cape_m import config +if args['debug']: + from configs.ae_body_cape_m import config os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] if args['debug']: diff --git a/scripts/cvae.py b/scripts/cvae.py index fcd36f0..87fff2d 100644 --- a/scripts/cvae.py +++ b/scripts/cvae.py @@ -12,10 +12,12 @@ 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"] = config['gpu'] -# os.environ["CUDA_LAUNCH_BLOCKING"] = '1' +if args['debug']: + os.environ["CUDA_LAUNCH_BLOCKING"] = '1' import time import argparse @@ -488,7 +490,7 @@ if __name__ == '__main__': step = load_model(model, config['pretrained_model']) if not args['debug']: - model = torch.compile(model) + # model = torch.compile(model) wandb.watch(model, log='all', log_freq=1000, log_graph=True) - + train() \ No newline at end of file diff --git a/scripts/upscaler_train.py b/scripts/upscaler_train.py index 7abf6f5..f4f6dfe 100644 --- a/scripts/upscaler_train.py +++ b/scripts/upscaler_train.py @@ -38,7 +38,7 @@ from utils.utils import InfiniteDataLoader as DataLoader from pytorch_msssim import MS_SSIM -low_res = 32 +low_res = 64 is_skirt = False is_cape = True is_male = True diff --git a/utils/checkpoints.py b/utils/checkpoints.py index 07ebcae..ebbd87a 100644 --- a/utils/checkpoints.py +++ b/utils/checkpoints.py @@ -31,7 +31,7 @@ class CheckpointSaver: self.top_model_paths = self.top_model_paths[len(self.top_model_paths)-self.top_n:] -def load_model(model, checkpoint, optimizer=None, scaler=None, iteration=''): +def load_model(model: torch.nn.Module, checkpoint, optimizer: torch.nn.Module=None, scaler: torch.nn.Module=None, iteration=''): # list of checkpoint files in the folder checkpoints checkpoints = [os.path.join(checkpoint, f) for f in os.listdir(checkpoint) if f.endswith(f'{iteration}.pt')] checkpoint = torch.load(checkpoints[-1])