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

more training

Browse Source
This commit is contained in:
luni
2023-02-25 22:46:15 +00:00
parent 90cf267dd8
commit aa96c6eb9c
6 changed files with 234 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
aaa_helpers.txt
+7
View File
@@ -14,6 +14,13 @@ 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_skirt" > nohup/iccv/upscale/iccv_32_skirt.out 2>&1 &
skirts!!!!
pyramid next week
auto_encoder_static.py
+2 -2
View File
@@ -247,9 +247,9 @@ if __name__ == '__main__':
load_model(model.decoder, config['pretrained_model'] + '_dec')
ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + config['name'] + '_enc', max_to_keep=20)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/iccv/' + config['name'] + '_enc', max_to_keep=20)
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/post/' + config['name'] + '_dec', max_to_keep=20)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/iccv/' + config['name'] + '_dec', max_to_keep=20)
# test()
train()
configs/iccv/config_abl_pyr_32.py
+184
View File
@@ -0,0 +1,184 @@
config = dict()
config['gpu'] = '2'
config['name'] = 'iccv_pyr_32'
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/local-data/train'
config['valid_dir'] = '/home/workspace/local-data/valid'
config['test_dir'] = '/home/workspace/data/dataset/test_ready'
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['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'] = True
config['pyramid']['noise_amount'] = 0.001
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/post/static_post_2'
config['pyramid']['body_model'] = 'checkpoints/post/body_post_2'
config['pyramid']['curvature_model'] = 'checkpoints/post/curv_post_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'] = config['latent_dim']
level['kernel_size'] = 7
level['encoder_depths'] = [3, 3, 9, 3]
level['encoder_dims'] = [32, 64, 128, 64]
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'] = config['latent_dim']
level['kernel_size'] = 7
level['encoder_depths'] = [9, 9, 9, 9]
level['encoder_dims'] = [32, 64, 128, 64]
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_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 128:
level = dict()
level['resolution'] = 128
level['latent_dim'] = config['latent_dim']
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['avg_garment_path'] = ''
level['upscale_enhancer'] = ''#'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['encoder_depths'] = [9, 9]
level['encoder_dims'] = [64, 64]
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_freeze'] = True
config['pyramid']['levels'].append(level)
elif res == 32:
level = dict()
level['resolution'] = 32
level['kernel_size'] = 3
level['encoder_depths'] = [16]
level['encoder_dims'] = [64]
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_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['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']['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['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['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'
configs/iccv/config_ae_static_cape_m.py
+2 -2
View File
@@ -33,14 +33,14 @@ config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims']
config['cond_ae']['activation'] = 'tanh'
config['gpu'] = '5'
config['gpu'] = '3'
config['lr'] = 0.0005
config['name'] = 'static_iccv_1_cape_m'
config['notes'] = ' '
config['tags'] = ['iccv', 'static', 'cape_m']
config['cond_ae']['upscale_output'] = False
config['pretrained_model'] = ''
config['pretrained_model'] = 'checkpoints/iccv/static_iccv_1_cape_m_old'
config['signature_dynamic'] = { 'id': (), 'frame_nr': (), }
models/upscaler.py
+4 -4
View File
@@ -131,13 +131,13 @@ class UpscalerBasicConv(tf.keras.Model):
self.low_res = low_res
self.seq = tf.keras.Sequential([
tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
tf.keras.layers.Conv2D(16, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
])
upscale_test.py
+32 -21
View File
@@ -1,5 +1,5 @@
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "6"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
@@ -35,7 +35,15 @@ from models.upscaler import UpscalerBasicConv as Upscaler
# !!!!!!!!!! CHANGE SSIM BASED ON RES
# !!!!!!!!!! CHANGE SKIRT OR NOT
name='256_512_skirt'
low_res = 64
is_skirt = True
is_cape = False
skirt_modifier = '_skirt' if is_skirt else ''
smpl_res = 3638 if is_cape else 14475
name=f'iccv_{low_res}_skirt'
def init_wandb():
run = wandb.init(project='upscale', name=name,
notes='\
@@ -59,15 +67,10 @@ smaller features | \
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
is_skirt = True
is_cape = False
skirt_modifier = '_skirt' if is_skirt else ''
smpl_res = 3638 if is_cape else 14475
uv2mesh = UV2Mesh(res=256, is_skirt=is_skirt, is_smpl_orig=is_cape)
def plot_3d(mesh_gt, mesh_pred_unposed, faces):
mesh_pred_unposed[:,2] += 1
fig = go.Figure(data=[
@@ -105,7 +108,7 @@ def plot_3d(mesh_gt, mesh_pred_unposed, faces):
return fig
# @tf.function
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=True)
@@ -115,8 +118,12 @@ def train_step(batch):
mesh_orig = batch['mesh_unposed']
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed' + high_resolution_modifier] - pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis]) )), axis=-1))
if low_res == 32:
ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2))
else:
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1)))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2))
loss = 1e3 * (1 - ms_ssim)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
@@ -125,7 +132,6 @@ def train_step(batch):
return { 'train/loss': loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim }
@tf.function
def valid_step(batch):
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False)
@@ -135,15 +141,18 @@ def valid_step(batch):
mesh_orig = batch['mesh_unposed']
l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_unposed' + high_resolution_modifier] - pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis]) )), axis=-1))
if low_res == 32:
ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2))
else:
ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2, power_factors=(0.5, 0.2, 0.2, 0.1)))
# ms_ssim = tf.reduce_mean(tf.image.ssim(pred * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, batch['uv_unposed' + high_resolution_modifier] * batch['uv_mask' + high_resolution_modifier][:,:,:,tf.newaxis] + 1, max_val=2))
loss = 1e3 * (1 - ms_ssim)
metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
return { 'valid/loss': loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim }
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_unposed' + low_resolution_modifier], batch['uv_mask' + low_resolution_modifier]], training=False)
@@ -168,7 +177,6 @@ def plot_everything(batch: dict, type: str):
return { **figs, **{ type + '/images': imgs} }
def train():
valid_iterator = iter(dataset_valid)
@@ -204,8 +212,8 @@ def train():
step += 1
if __name__ == '__main__':
# init_wandb()
allow_memory_growth()
init_wandb()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())
@@ -214,7 +222,6 @@ if __name__ == '__main__':
# uv_mask_body = np.load('custom_utils/mask_256x512.npy')
# uv_mask_body = cv2.dilate(uv_mask_body.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2)
low_res = 32
high_res = 512
uv_mask_body = tf.ones((low_res, low_res//2))
low_resolution_modifier = '_' + str(low_res)
@@ -234,8 +241,8 @@ if __name__ == '__main__':
'uv_mask' + low_resolution_modifier: (low_res, low_res//2,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train', 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/train', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
dataset_train = factory.get_dataset(folder='/home/workspace/data/dataset/local-cache/train_skirt', 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/data/dataset/local-cache/train_skirt', 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/valid_skirt', 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/valid_skirt', signature=signature_static, resolutions=[high_res, low_res], in_memory=True, is_cape=is_cape)
@@ -244,14 +251,18 @@ if __name__ == '__main__':
# 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)
if low_res == 32 or low_res == 64:
model = Upscaler(low_res=low_res, kernel_size=3)
elif low_res == 128 or low_res == 256:
model = Upscaler(low_res=low_res, kernel_size=5)
else:
model = Upscaler(low_res=low_res, kernel_size=7)
optimizer = tf.optimizers.Adam(learning_rate=0.0001)
folder = 'models/ae_logs/'
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=20)
load_model(model, 'checkpoints/32_512_4')
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/iccv/' + name, max_to_keep=20)
train()