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

started training pyramid

Browse Source
This commit is contained in:
luni
2023-02-26 22:25:50 +00:00
parent a678724711
commit 64ecb900f5
8 changed files with 314 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aaa_helpers.txt
+20 -5
View File
@@ -14,12 +14,27 @@ nohup python auto_encoder_curv.py --name="lhunor_curv_iccv_1_cape_m" > nohup/icc
nohup python cvae.py --name="lhunor_abl_all" > nohup/iccv/ablation/abl_all.out 2>&1 &
nohup python upscale_test.py --name="iccv_32" > nohup/iccv/upscale/iccv_32.out 2>&1 &
nohup python upscale_test.py --name="iccv_64" > nohup/iccv/upscale/iccv_64.out 2>&1 &
nohup python upscale_test.py --name="iccv_128" > nohup/iccv/upscale/iccv_128.out 2>&1 &
nohup python upscale_test.py --name="iccv_256" > nohup/iccv/upscale/iccv_256.out 2>&1 &
nohup python upscale_test.py --name="iccv_32_2" > nohup/iccv/upscale/iccv_32_2.out 2>&1 &
nohup python upscale_test.py --name="iccv_64_2" > nohup/iccv/upscale/iccv_64_2.out 2>&1 &
nohup python upscale_test.py --name="iccv_128_2" > nohup/iccv/upscale/iccv_128_2.out 2>&1 &
nohup python upscale_test.py --name="iccv_256_2" > nohup/iccv/upscale/iccv_256_2.out 2>&1 &
nohup python upscale_test.py --name="iccv_32_skirt" > nohup/iccv/upscale/iccv_32_skirt.out 2>&1 &
nohup python upscale_test.py --name="iccv_32_2_skirt" > nohup/iccv/upscale/iccv_32_2_skirt.out 2>&1 &
nohup python upscale_test.py --name="iccv_64_2_skirt" > nohup/iccv/upscale/iccv_64_2_skirt.out 2>&1 &
nohup python upscale_test.py --name="iccv_128_2_skirt" > nohup/iccv/upscale/iccv_128_2_skirt.out 2>&1 &
nohup python upscale_test.py --name="iccv_256_2_skirt" > nohup/iccv/upscale/iccv_256_2_skirt.out 2>&1 &
nohup python upscale_test.py --name="iccv_32_2_cape_f" > nohup/iccv/upscale/iccv_32_2_cape_f_real.out 2>&1 &
nohup python upscale_test.py --name="iccv_64_2_cape_f" > nohup/iccv/upscale/iccv_64_2_cape_f.out 2>&1 &
nohup python upscale_test.py --name="iccv_128_2_cape_f" > nohup/iccv/upscale/iccv_128_2_cape_f.out 2>&1 &
nohup python upscale_test.py --name="iccv_256_2_cape_f" > nohup/iccv/upscale/iccv_256_2_cape_f.out 2>&1 &
nohup python upscale_test.py --name="iccv_32_2_cape_m" > nohup/iccv/upscale/iccv_32_2_cape_m.out 2>&1 &
nohup python upscale_test.py --name="iccv_64_2_cape_m" > nohup/iccv/upscale/iccv_64_2_cape_m.out 2>&1 &
nohup python upscale_test.py --name="iccv_128_2_cape_m" > nohup/iccv/upscale/iccv_128_2_cape_m.out 2>&1 &
nohup python upscale_test.py --name="iccv_256_2_cape_m" > nohup/iccv/upscale/iccv_256_2_cape_m.out 2>&1 &
nohup python cvae_pyramid.py --name="lhunor_iccv_pyr_32" > nohup/iccv/pyr/iccv_pyr_32.out 2>&1 &
skirts!!!!
pyramid next week
configs/iccv/config_abl_pyr_32.py
+29 -25
View File
@@ -1,6 +1,6 @@
config = dict()
config['gpu'] = '2'
config['gpu'] = '0'
config['name'] = 'iccv_pyr_32'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv']
@@ -47,6 +47,7 @@ config['lr'] = 0.0001
config['starting_balancing_coeff'] = 0.001
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
@@ -63,8 +64,8 @@ config['signature_static'] = { 'uv_static': (512, 256, 3),
}
config['pyramid']['is_label'] = False
config['pyramid']['add_noise'] = True
config['pyramid']['noise_amount'] = 0.001
config['pyramid']['add_noise'] = False
config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
@@ -77,9 +78,10 @@ for res in config['pyramid']['resolutions']:
config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
config['pyramid']['static_model'] = 'checkpoints/post/static_post_2'
config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_1'
config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
config['pyramid']['freeze_encoders'] = True
@@ -88,10 +90,10 @@ for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['latent_dim'] = config['latent_dim']
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['encoder_depths'] = [3, 3, 9, 3]
level['encoder_dims'] = [32, 64, 128, 64]
level['encoder_dims'] = [32, 64, 128, level['latent_dim']]
level['decoder_input'] = (16, 8, 80)
level['decoder_depths'] = [3, 3, 9, 3]
level['decoder_dims'] = [256, 128, 64, 16]
@@ -101,53 +103,55 @@ for res in config['pyramid']['resolutions']:
elif res == 256:
level = dict()
level['resolution'] = 256
level['latent_dim'] = config['latent_dim']
level['kernel_size'] = 7
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [9, 9, 9, 9]
level['encoder_dims'] = [32, 64, 128, 64]
level['encoder_dims'] = [32, 64, 128, level['latent_dim']]
level['decoder_input'] = (8, 4, 128)
level['decoder_depths'] = [9, 9, 9, 9]
level['decoder_dims'] = [256, 128, 64, 16]
level['avg_garment_path'] = ''
level['upscale_enhancer'] = ''#'checkpoints/up_3_d256_512_2'
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d256_512_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['latent_dim'] = config['latent_dim']
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [16, 9, 9, 9]
level['encoder_dims'] = [64, 128, 128, 64]
level['decoder_input'] = (4, 2, 128)
level['decoder_depths'] = [16, 9, 9, 9]
level['decoder_dims'] = [256, 128, 128, 64]
level['decoder_dims'] = [256, 128, 128, level['latent_dim']]
level['avg_garment_path'] = ''
level['upscale_enhancer'] = ''#'checkpoints/up_3_d128_256_2'
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d128_256_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['kernel_size'] = 5
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [9, 9]
level['encoder_dims'] = [64, 64]
level['encoder_dims'] = [64, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/64_512_2'#checkpoints/up_3_d64_128_2'
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#checkpoints/up_3_d64_128_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [16]
level['encoder_dims'] = [64]
level['encoder_dims'] = [level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/32_512_4'#'checkpoints/up_3_d32_64_2'
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
else:
@@ -156,9 +160,9 @@ for res in config['pyramid']['resolutions']:
config['cond_ae_static'] = dict()
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8
config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 9, 3]
config['cond_ae_static']['encoder_dims'] = [16, 32, 64, 64, config['cond_ae_static']['bottleneck_size']]
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
config['cond_ae_static']['activation'] = 'tanh'
@@ -166,7 +170,7 @@ config['cond_ae_static']['activation'] = 'tanh'
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
@@ -176,7 +180,7 @@ config['cond_ae_curv']['activation'] = 'tanh'
config['cond_ae_body'] = dict()
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
configs/iccv/config_abl_pyr_32_skirt.py
+188
View File
@@ -0,0 +1,188 @@
config = dict()
config['gpu'] = '0'
config['name'] = 'iccv_pyr_32_skirt'
config['notes'] = ''
config['tags'] = ['pyr', 'iccv']
# DON'T FORGET TO RESET SSIM LOSS USAGE
# DON'T FORGET TO CHECK MODEL LOADING
# DON'T FORGET TO CHECK IF STORE IN MEMORY
# DON'T FORGET TO CHECK REG LOSS COEFF
# DON'T FORGET TO LOAD FINAL UPSCALER UNLESS HAVE LEVEL 512
# DON'T FORGET TO CHECK IF SKIRT
# DON'T FORGET TO CHECK IF CAPE
# DON'T FORGET TO CHECK IF USING DOWNSCALE NONE OR _2
# DON'T FORGET TO CHECK UPSCALING and ENCODER CHECKPOINTS
config['use_skinning'] = True # and deshaping
config['remove_global_orientation'] = True
config['use_extra_losses'] = True # edge, collision, 3d mesh
config['use_offsets'] = False
config['use_norm_loss'] = True
config['pretrained_model'] = ''
config['pretrained_model_iteration'] = ''
config['pyramid'] = dict()
config['pyramid']['resolutions'] = [32]
config['pyramid']['pretrained_pyramid'] = ''
config['pyramid']['final_upscale_enhancer'] = ''
config['train_dir'] = '/home/workspace/data/dataset/local-cache/train_skirt'
config['valid_dir'] = '/home/workspace/local-data/valid_skirt'
config['test_dir'] = '/home/workspace/data/dataset/test_skirt_ready_2'
config['batch_size'] = 4
config['is_skirt'] = False
config['is_cape'] = False
config['dataset_in_memory'] = False
config['resolution'] = None
config['smpl_res'] = 14475
if config['is_cape']:
config['smpl_res'] = 3638
config['lr'] = 0.0001
config['starting_balancing_coeff'] = 0.001
config['reg_loss_coeff'] = 0.1 # set to 0 if don't want to use edge+collision
config['avg_garment_path'] = ''
config['bottleneck_scale'] = 2
config['signature_dynamic'] = { 'id': (), 'frame_nr': (),
'pose_shape': (72 + 11,),
'uv_body_posed': (512, 256, 3),
'mesh_body_posed': (config['smpl_res'], 3,),
'mesh_posed': (config['smpl_res'], 3,),
'mesh_unposed': (config['smpl_res'], 3,),
'uv_unposed': (512, 256, 3),
'uv_curvature': (512, 256, 7),
}
config['signature_static'] = { 'uv_static': (512, 256, 3),
'mesh_mask': (config['smpl_res'],),
'uv_mask': (512, 256,),
}
config['pyramid']['is_label'] = False
config['pyramid']['add_noise'] = False
config['pyramid']['noise_amount'] = 0.0
config['pyramid']['loss_weights'] = [1,1,1,1,1]
for res in config['pyramid']['resolutions']:
if res != 512:
res_modifier = '_' + str(res)
config['signature_dynamic']['uv_body_posed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_unposed' + res_modifier] = (res, res//2, 3)
config['signature_dynamic']['uv_curvature' + res_modifier] = (res, res//2, 7)
config['signature_static']['uv_static' + res_modifier] = (res, res//2, 3)
config['signature_static']['uv_mask' + res_modifier] = (res, res//2)
config['pyramid']['static_model'] = 'checkpoints/iccv/static_iccv_1'
config['pyramid']['body_model'] = 'checkpoints/iccv/body_iccv_1'
config['pyramid']['curvature_model'] = 'checkpoints/iccv/curv_iccv_1'
config['pyramid']['freeze_encoders'] = True
config['pyramid']['levels'] = []
for res in config['pyramid']['resolutions']:
if res == 512:
level = dict()
level['resolution'] = 512
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 7
level['encoder_depths'] = [3, 3, 9, 3]
level['encoder_dims'] = [32, 64, 128, level['latent_dim']]
level['decoder_input'] = (16, 8, 80)
level['decoder_depths'] = [3, 3, 9, 3]
level['decoder_dims'] = [256, 128, 64, 16]
level['avg_garment_path'] = ''
level['upscale_enhancer'] = ''
config['pyramid']['levels'].append(level)
elif res == 256:
level = dict()
level['resolution'] = 256
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [9, 9, 9, 9]
level['encoder_dims'] = [32, 64, 128, level['latent_dim']]
level['decoder_input'] = (8, 4, 128)
level['decoder_depths'] = [9, 9, 9, 9]
level['decoder_dims'] = [256, 128, 64, 16]
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d256_512_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 5
level['encoder_depths'] = [16, 9, 9, 9]
level['encoder_dims'] = [64, 128, 128, 64]
level['decoder_input'] = (4, 2, 128)
level['decoder_depths'] = [16, 9, 9, 9]
level['decoder_dims'] = [256, 128, 128, level['latent_dim']]
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#'checkpoints/up_3_d128_256_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 64:
level = dict()
level['resolution'] = 64
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [9, 9]
level['encoder_dims'] = [64, level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_'#checkpoints/up_3_d64_128_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['latent_dim'] = 64 * config['bottleneck_scale']
level['kernel_size'] = 3
level['encoder_depths'] = [16]
level['encoder_dims'] = [level['latent_dim']]
level['decoder_depths'] = level['encoder_depths']
level['decoder_dims'] = level['encoder_dims']
level['avg_garment_path'] = ''
level['upscale_enhancer'] = 'checkpoints/iccv/iccv_32_2'#'checkpoints/up_3_d32_64_2'
level['upscale_enhancer_freeze'] = True
config['pyramid']['levels'].append(level)
else:
raise NotImplementedError
config['cond_ae_static'] = dict()
config['cond_ae_static']['is_label'] = False
config['cond_ae_static']['kernel_size'] = 7
config['cond_ae_static']['bottleneck_size'] = 8 * config['bottleneck_scale']
config['cond_ae_static']['encoder_depths'] = [3, 3, 3, 27, 3]
config['cond_ae_static']['encoder_dims'] = [8, 16, 32, 64, config['cond_ae_static']['bottleneck_size']]
config['cond_ae_static']['decoder_depths'] = config['cond_ae_static']['encoder_depths']
config['cond_ae_static']['decoder_dims'] = config['cond_ae_static']['encoder_dims']
config['cond_ae_static']['activation'] = 'tanh'
config['cond_ae_curv'] = dict()
config['cond_ae_curv']['is_label'] = True
config['cond_ae_curv']['kernel_size'] = 7
config['cond_ae_curv']['bottleneck_size'] = 40 * config['bottleneck_scale']
config['cond_ae_curv']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_curv']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_curv']['bottleneck_size']]
config['cond_ae_curv']['decoder_depths'] = config['cond_ae_curv']['encoder_depths']
config['cond_ae_curv']['decoder_dims'] = config['cond_ae_curv']['encoder_dims']
config['cond_ae_curv']['activation'] = 'tanh'
config['cond_ae_body'] = dict()
config['cond_ae_body']['is_label'] = False
config['cond_ae_body']['kernel_size'] = 7
config['cond_ae_body']['bottleneck_size'] = 16 * config['bottleneck_scale']
config['cond_ae_body']['encoder_depths'] = [3, 3, 9, 9, 3]
config['cond_ae_body']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae_body']['bottleneck_size']]
config['cond_ae_body']['decoder_depths'] = config['cond_ae_body']['encoder_depths']
config['cond_ae_body']['decoder_dims'] = config['cond_ae_body']['encoder_dims']
config['cond_ae_body']['activation'] = 'tanh'
cvae.py
+28 -12
View File
@@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes
from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.cvae_model import CVAENew as CVAE
from custom_utils.custom_utils import Accumulator
def init_wandb():
@@ -529,7 +530,8 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + config['name'], max_to_keep=20)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
if config['pretrained_model'] != '':
load_model(model, config['pretrained_model'])
@@ -542,31 +544,45 @@ if __name__ == '__main__':
batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic']))
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 10 == 0:
if step % 50 == 0:
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
# to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
# if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
# train_figs = plot_everything(batch, 'train')
# valid_figs = plot_everything(data_valid, 'valid')
# to_log = { **to_log, **train_figs, **valid_figs }
if step % 5000 == 0:
manager.save(checkpoint_number=step)
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
# if step % 5000 == 0:
# manager.save(checkpoint_number=step)
# if step == 0:
# to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < (config['balancing_activation_limit'] if 'balancing_activation_limit' in config.keys() else 0.017):
if balancing_activated_step == 0:
balancing_activated_step = step
g_balancing_coeff = g_balancing_coeff + 1 / 10000
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
to_log = { **to_log, **{ 'train/balancing_coeff': g_balancing_coeff } }
wandb.log(to_log, step = step)
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
cvae_pyramid.py
+24 -16
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.post.config_abl_pyr_64_128_5 import config
from configs.iccv.config_abl_pyr_32 import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -43,10 +43,11 @@ from dataset import DatasetFactory, StaticDataset
from custom_utils.smpl import SMPL
from models.pyramid_model_optimized import PyramidNew as Pyramid
from postprocessing import fix_collision
from custom_utils.custom_utils import Accumulator
def init_wandb():
run = wandb.init(project='ablation', name=config['name'], notes=config['notes'], dir='/wandb', tags=config['tags'])
run = wandb.init(project='iccv', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('cvae_pyramid.py')
@@ -61,11 +62,14 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py')
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.post.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
file_name = '/'.join(config_file_name.split('.')) + '.py'
if os.path.exists(file_name):
artifact.add_file(file_name)
run.log_artifact(artifact)
uv2mesh = UV2Mesh(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
def plot_3d(mesh_gt, mesh_pred_unposed_levels, faces, mesh_posed, collision, collision_faces, collision_color):
fig = make_subplots(rows=1, cols=1+len(mesh_pred_unposed_levels)+1+1, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(1+len(mesh_pred_unposed_levels)+1+1)] ])
@@ -791,6 +795,8 @@ if __name__ == '__main__':
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/post/' + config['name'], max_to_keep=20)
train_accumulator = Accumulator()
valid_accumulator = Accumulator()
# load_model(model, 'checkpoints/pyr_abl_32_skirt', iteration='180000')
@@ -846,31 +852,33 @@ if __name__ == '__main__':
train_loss = train_step(batch)
train_accumulator.update(train_loss)
if step % 50 == 0:
data_valid = valid_iterator.get_next()
data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
valid_loss = valid_step(data_valid)
valid_accumulator.update(valid_loss)
to_log = { **train_loss, **valid_loss, **{ 'current_step': step} }
if step % 5000 == 0 or (step < 5000 and step % 500 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step % 5000 == 0:
manager.save(checkpoint_number=step)
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
if train_loss['train/euclid_error'] < 0.004:
if train_loss['train/euclid_error'] < 0.015:
if balancing_activated_step == 0:
balancing_activated_step = step
balancing_coeff = balancing_coeff + 1 / 20000
if step % 1000 == 0 or (step < 1000 and step % 100 == 0):
to_log = dict({ **train_loss, **valid_loss, **{ 'current_step': step } })
if step % 5000 == 0 or (step < 5000 and step % 1000 == 0):
train_figs = plot_everything(batch, 'train')
valid_figs = plot_everything(data_valid, 'valid')
to_log = { **to_log, **train_figs, **valid_figs }
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
to_log = { **to_log, **{ 'train/balancing_coeff': balancing_coeff, 'train/gamma': gamma } }
wandb.log(to_log, step = step)
wandb.log(to_log, step=step)
train_accumulator.reset()
valid_accumulator.reset()
print(step)
models/pyramid_level.py
+3 -4
View File
@@ -102,13 +102,12 @@ class PyramidLevelNew(tf.keras.Model):
super(PyramidLevelNew, self).__init__()
self.size = config['resolution']
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'])
self.encoder = VariationalEncoderNew(config['encoder_depths'], config['encoder_dims'], kernel_size=config['kernel_size'], latent_dim=config['latent_dim'])
self.decoder = DecoderNew(config['decoder_depths'], config['decoder_dims'], config['kernel_size'])
self.upscale_enhancer = None
if config['upscale_enhancer'] != '':
k_size = config['upscale_enhancer_kernel_size'] if 'upscale_enhancer_kernel_size' in config.keys() else ((3 if config['resolution'] <= 64 else 5))
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=k_size)
self.upscale_enhancer = UpscalerBasicConv(config['resolution'], kernel_size=config['kernel_size'])
load_model(self.upscale_enhancer, config['upscale_enhancer'])
self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
@@ -117,7 +116,7 @@ class PyramidLevelNew(tf.keras.Model):
else:
self.avg_garment = None
self.reshape = tf.keras.layers.Reshape((8, 4, 64))
self.reshape = tf.keras.layers.Reshape((8, 4, config['latent_dim']))
def enhance(self, low, mask, training=None):
models/pyramid_model_optimized.py
-4
View File
@@ -343,11 +343,7 @@ class PyramidNew(tf.keras.Model):
for i in range(1, len(self.levels)):
input = levels_input[i]
input[0] = input[0] - pred
input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3])
input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None]
input[0] = input[0] * input[-2][:,:,:,None] # uv_mask_res
pred_level, mean, logvar = self.levels[i](input, z_cond, training=training)
pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res
pred = pred + pred_level
if i != len(self.levels) - 1:
pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)
upscale_test.py
+12 -3
View File
@@ -36,14 +36,17 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SKIRT OR NOT
low_res = 128
low_res = 256
is_skirt = True
is_cape = False
is_male = False
skirt_modifier = '_skirt' if is_skirt else ''
cape_modifier = '_cape' if is_cape else ''
gender_modifier = '' if not is_cape else ('_m' if is_male else '_f')
smpl_res = 3638 if is_cape else 14475
name=f'iccv_{low_res}_2' + skirt_modifier
name=f'iccv_{low_res}_2' + skirt_modifier + cape_modifier + gender_modifier
def init_wandb():
run = wandb.init(project='upscale', name=name,
notes='\
@@ -202,7 +205,6 @@ def train():
if step % 5000 == 0:
manager.save(checkpoint_number=step)
wandb.log(to_log, step = step)
print(step)
@@ -255,10 +257,17 @@ if __name__ == '__main__':
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if is_cape:
if is_male:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/male', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/male', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
else:
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/train/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/cape_processed/train/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/cape_processed/test/female', batch_size=8, signature=signature_dynamic, resolutions=[high_res, low_res], in_memory=False, is_cape=is_cape)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/cape_processed/test/female', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
if low_res == 32 or low_res == 64:
model = Upscaler(low_res=low_res, kernel_size=3)