diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 6a39d1e..11f8808 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -1,4 +1,5 @@ nohup python cvae_pyramid.py --name="lhunor_pyr_abl_32_64_e2e_r3" > nohup/review/pyr_abl_32_64_e2e_r3.out 2>&1 & +nohup python cvae_pyramid.py --name="lhunor_pyr_abl_32_64_e2e_r3_pre" > nohup/review/pyr_abl_32_64_e2e_r3_pre.out 2>&1 & nohup python evaluate_ae_pyramid.py --name evaluate_pyr_abl_32_64_e2e_r3 > nohup/evaluate/pyr_abl_32_64_e2e_r3.out 2>&1 & nohup python calculate_error.py --name error_pyr_abl_32_64_e2e_r3 > nohup/error/pyr_abl_32_64_e2e_r3.out 2>&1 & diff --git a/configs/config_abl_pyr_32.py b/configs/config_abl_pyr_32.py index a6a7060..5aeef36 100644 --- a/configs/config_abl_pyr_32.py +++ b/configs/config_abl_pyr_32.py @@ -1,7 +1,7 @@ config = dict() config['gpu'] = '0' -config['name'] = 'pyr_abl_32_64_e2e_r3' +config['name'] = 'pyr_abl_32_64_e2e_r3_pre' config['notes'] = '' config['tags'] = ['review_3'] @@ -26,7 +26,7 @@ config['pyramid']['resolutions'] = [32, 64] config['pyramid']['pretrained_pyramid'] = '' config['pyramid']['final_upscale_enhancer'] = 'checkpoints/64_512_2' -config['pretrained_model'] = '' +config['pretrained_model'] = 'checkpoints/pyr_abl_32_64_e2e_r3' config['pretrained_model_iteration'] = '' config['train_dir'] = '/home/workspace/local-data/train' diff --git a/cvae_pyramid.py b/cvae_pyramid.py index 2686cea..ffcb62c 100644 --- a/cvae_pyramid.py +++ b/cvae_pyramid.py @@ -807,9 +807,8 @@ if __name__ == '__main__': # assert len(config['pyramid']['resolutions']) == 1 load_model(model, config['pretrained_model'], iteration=config['pretrained_model_iteration']) - model.static_enc.trainable = False - model.body_enc.trainable = False - model.curv_enc.trainable = False + model.trainable = False + model.cond_encoder_2.trainable = True # if len(config['pyramid']['resolutions']) == 1: # else: # load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else '')) diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 9cf7322..4b847a1 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -92,7 +92,7 @@ class PyramidLevel(tf.keras.Model): # x = x * uv_mask_high[:,:,:,tf.newaxis] - return x, mean, logvar + return x, z, mean, logvar diff --git a/models/pyramid_model_optimized.py b/models/pyramid_model_optimized.py index 87eed8c..800d179 100644 --- a/models/pyramid_model_optimized.py +++ b/models/pyramid_model_optimized.py @@ -78,6 +78,11 @@ class Pyramid(tf.keras.Model): load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) self.upscale_enhancer.trainable = False + self.cond_encoder_2 = tf.keras.Sequential([tf.keras.layers.Dense(512, activation='relu'), + tf.keras.layers.Dense(256, activation='relu'), + tf.keras.layers.Dense(128, activation='relu'), + tf.keras.layers.Dense(config['pyramid']['latent_dim'], activation='tanh'),]) + def enhance(self, low, mask, training): x = self.upscale_enhancer([low, mask], training=training) return x @@ -90,19 +95,25 @@ class Pyramid(tf.keras.Model): z_cond = tf.concat([z_static, z_body, z_curv], axis=-1) z_cond = self.cond_encoder(z_cond, training=training) + return z_cond, z_static, z_body, z_curv + + def encode_conds_2(self, z, z_static, z_body, z_curv, training): + z_cond = tf.concat([z, z_static, z_body, z_curv], axis=-1) + z_cond = self.cond_encoder_2(z_cond, training=training) + return z_cond def call(self, inputs, training): levels_input, conds = inputs uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2] - z_cond = self.encode_conds(uv_static, uv_body, uv_curv, training) + z_cond, z_static, z_body, z_curv = self.encode_conds(uv_static, uv_body, uv_curv, training) if self.add_noise == True and training == True: z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) out_levels = [] - pred, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training) + pred, z, mean, logvar = self.levels[0](levels_input[0], z_cond, training=training) if self.use_offsets: pred = pred + levels_input[0][2] # uv_body_posed_32 if 1 != len(self.levels): @@ -119,8 +130,9 @@ class Pyramid(tf.keras.Model): for i in range(1, len(self.levels)): input = levels_input[i] input[0] = pred # use previous level output as input for next level + z_cond_2 = self.encode_conds_2(z, z_static, z_body, z_curv, training=training) - pred_level, mean, logvar = self.levels[i](input, z_cond, training=training) + pred_level, z, mean, logvar = self.levels[i](input, z_cond_2, training=training) pred = pred + pred_level if i != len(self.levels) - 1: pred = self.levels[i].enhance(pred, levels_input[i][-2], training=training)