diff --git a/aaa_helpers.txt b/aaa_helpers.txt index 99edec9..ca3d638 100644 --- a/aaa_helpers.txt +++ b/aaa_helpers.txt @@ -1,7 +1,7 @@ nohup python cvae.py --name="lhunor_weights_test_6" > nohup/iccv/weights_test_6.out 2>&1 & nohup python evaluate_ae.py --name iccv_weights_test_6 > nohup/iccv/evaluate_iccv_weights_test_6.out 2>&1 & nohup python calculate_error.py --name iccv_weights_test_6 > nohup/iccv/error_iccv_weights_test_6.out 2>&1 & -nohup python auto_encoder_static.py --name="lhunor_ae_static_iccv_resnet_1" > nohup/iccv/ae_static_iccv_resnet_1.out 2>&1 & +nohup python auto_encoder_static.py --name="lhunor_ae_static_iccv_resnet_1" > nohup/iccv/encoders/ae_static_iccv_resnet_1.out 2>&1 & one final convolutional layer for smoothing, on the final coordinates diff --git a/auto_encoder_body.py b/auto_encoder_body.py index 83eb491..e531bf2 100644 --- a/auto_encoder_body.py +++ b/auto_encoder_body.py @@ -1,6 +1,6 @@ import os import sys -from configs.review.config_ae_body import config +from configs.iccv.config_ae_body import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -31,10 +31,11 @@ from losses import edge_loss, uv_mask_loss from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, plot_3d from models.cond_encoder import AutoEncoderNew as AutoEncoder +from custom_utils.custom_utils import Accumulator def init_wandb(): - run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags']) + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_body.py') artifact.add_file('losses.py') @@ -42,9 +43,11 @@ def init_wandb(): artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') - config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.review.')] + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] 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_smpl_orig=config['is_cape'], is_skirt=config['is_skirt']) @@ -60,7 +63,8 @@ def train_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475)))) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=3)) loss = 1e3 * (1 - ms_ssim) @@ -80,7 +84,8 @@ def valid_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, tf.ones((pred.shape[0], 3638 if config['is_cape'] else 14475)))) l1_loss = tf.reduce_mean(tf.reduce_sum(((tf.math.abs(batch['uv_body_posed' + resolution_modifier] - pred))), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_body_posed' + resolution_modifier] + 1, max_val=3)) loss = 1e3 * (1 - ms_ssim) @@ -124,24 +129,27 @@ def train(): for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) + train_accumulator.update(train_loss) - 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) + 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 == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } + 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)) } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() - if step % 1000 == 0: - wandb.log(to_log, step=step, commit=True) - else: - wandb.log(to_log, step=step, commit=False) - if step % 10000 == 0: manager_enc.save(checkpoint_number=step) manager_dec.save(checkpoint_number=step) @@ -207,5 +215,7 @@ if __name__ == '__main__': manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + 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) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() train() \ No newline at end of file diff --git a/auto_encoder_curv.py b/auto_encoder_curv.py index 7e75b4f..a891e4f 100644 --- a/auto_encoder_curv.py +++ b/auto_encoder_curv.py @@ -1,6 +1,6 @@ import os import sys -from configs.review.config_ae_curv import config +from configs.iccv.config_ae_curv import config os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu'] import tensorflow as tf @@ -32,12 +32,13 @@ from tf_utils import check_tf_version, allow_memory_growth from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict from models.cond_encoder import AutoEncoderNew as AutoEncoder +from custom_utils.custom_utils import Accumulator testing = False def init_wandb(): if not testing: - run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags']) + run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/tmp', tags=config['tags']) artifact = wandb.Artifact('code', type='code') artifact.add_file('auto_encoder_curv.py') artifact.add_file('losses.py') @@ -45,9 +46,11 @@ def init_wandb(): artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') - config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.review.')] + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] 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 = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) @@ -164,24 +167,26 @@ def train(): batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) + train_accumulator.update(train_loss) - 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) + 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 } - 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 == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - - if step % 1000 == 0: - wandb.log(to_log, step=step, commit=True) - else: - wandb.log(to_log, step=step, commit=False) - # print(step) + 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)) } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() if step % 10000 == 0: manager_enc.save(checkpoint_number=step) @@ -257,6 +262,8 @@ if __name__ == '__main__': manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/post/' + 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) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() if testing: test() diff --git a/auto_encoder_static.py b/auto_encoder_static.py index be8e02e..fc90622 100644 --- a/auto_encoder_static.py +++ b/auto_encoder_static.py @@ -34,6 +34,7 @@ from tf_utils import check_tf_version, allow_memory_growth from dataset import DatasetFactory, StaticDataset from custom_utils.custom_utils import batch_to_dict, plot_3d from models.cond_encoder import AutoEncoderNew as AutoEncoder +from custom_utils.custom_utils import Accumulator def init_wandb(): @@ -46,9 +47,11 @@ def init_wandb(): artifact.add_file('models/cond_encoder.py') artifact.add_file('dataset.py') artifact.add_file('convnext/models/convnext_tf.py') - config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.review.')] + config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.')] 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']) @@ -65,11 +68,14 @@ def train_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_static'] + 1, max_val=3)) loss = 1e3 * (1 - ms_ssim) + # loss = l1_loss 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']) gradients = tape.gradient(loss, model.trainable_variables) + gradients = [(tf.clip_by_norm(grad, clip_norm=2.0)) for grad in gradients] optimizer.apply_gradients(zip(gradients, model.trainable_variables)) return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/l1_loss': l1_loss, 'train/euclid_error': metric, 'train/ms_ssim': ms_ssim, } @@ -84,8 +90,10 @@ def valid_step(batch): e_loss = 1e-2 * tf.reduce_mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) l1_loss = tf.reduce_mean(tf.reduce_sum((tf.math.abs(batch['uv_static'] - pred)), axis=1)) - ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=2, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - loss = 1e3 * (1 - ms_ssim) + # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 1, batch['uv_static'] + 1, max_val=3, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) + ms_ssim = tf.reduce_mean(tf.image.ssim(pred + 1, batch['uv_static'] + 1, max_val=3)) + loss = 1e3 * (1 - ms_ssim) + # loss = l1_loss 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/edge_loss': e_loss, 'valid/l1_loss': l1_loss, 'valid/euclid_error': metric, 'valid/ms_ssim': ms_ssim, } @@ -126,21 +134,26 @@ def train(): for batch in dataset_train: batch = static_dataset_train.expand_batch(batch_to_dict(batch, config['signature_dynamic'])) train_loss = train_step(batch) + train_accumulator.update(train_loss) - 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) + 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 = dict({ **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 == 0: - to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } } - - - wandb.log(to_log, step=step) + 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)) } } + + wandb.log(to_log, step=step) + train_accumulator.reset() + valid_accumulator.reset() if step % 10000 == 0: manager_enc.save(checkpoint_number=step) @@ -184,17 +197,6 @@ def test(): print() - -class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule): - - def __init__(self, initial_learning_rate): - self.initial_learning_rate = initial_learning_rate - - def __call__(self, step): - if step < 1000: - return self.initial_learning_rate / 10.0 - else: - return self.initial_learning_rate if __name__ == '__main__': @@ -242,6 +244,8 @@ if __name__ == '__main__': model : AutoEncoder = AutoEncoder(config['cond_ae']) optimizer = tf.optimizers.Adam(learning_rate=config['lr']) + train_accumulator = Accumulator() + valid_accumulator = Accumulator() if config['pretrained_model'] != '': load_model(model.encoder, config['pretrained_model'] + '_enc') diff --git a/configs/iccv/config_ae_body.py b/configs/iccv/config_ae_body.py index 5018050..308abc3 100644 --- a/configs/iccv/config_ae_body.py +++ b/configs/iccv/config_ae_body.py @@ -23,19 +23,16 @@ config['cond_ae']['is_label'] = False config['cond_ae']['upscale_output'] = False config['cond_ae']['latent_dim'] = config['latent_dim'] config['cond_ae']['kernel_size'] = 7 -config['cond_ae']['bottleneck_size'] = 16 -config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] -config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']] -config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths'] -config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] +config['cond_ae']['bottleneck_size'] = 16 * 2 +config['cond_ae']['layer_params'] = [3, 7, 15, 7, 3] config['cond_ae']['activation'] = 'tanh' -config['gpu'] = '4' +config['gpu'] = '3' config['lr'] = 0.0001 -config['name'] = 'body_r7_3' -config['notes'] = 'resnet18' -config['tags'] = ['review_7'] +config['name'] = 'body_resnet_1' +config['notes'] = '' +config['tags'] = ['iccv', 'resnet'] config['cond_ae']['upscale_output'] = False config['pretrained_model'] = '' diff --git a/configs/iccv/config_ae_curv.py b/configs/iccv/config_ae_curv.py index 3a54856..a3a2b0c 100644 --- a/configs/iccv/config_ae_curv.py +++ b/configs/iccv/config_ae_curv.py @@ -23,19 +23,17 @@ config['cond_ae']['is_label'] = True config['cond_ae']['upscale_output'] = False config['cond_ae']['latent_dim'] = config['latent_dim'] config['cond_ae']['kernel_size'] = 7 -config['cond_ae']['bottleneck_size'] = 40 -config['cond_ae']['encoder_depths'] = [3, 3, 9, 9, 3] -config['cond_ae']['encoder_dims'] = [16, 32, 64, 128, config['cond_ae']['bottleneck_size']] -config['cond_ae']['decoder_depths'] = config['cond_ae']['encoder_depths'] -config['cond_ae']['decoder_dims'] = config['cond_ae']['encoder_dims'] +config['cond_ae']['bottleneck_size'] = 40 * 2 +config['cond_ae']['layer_params'] = [3, 7, 15, 7, 3] config['cond_ae']['activation'] = 'tanh' -config['gpu'] = '3' + +config['gpu'] = '2' config['lr'] = 0.001 -config['name'] = 'curv_r7_3' -config['notes'] = 'resnet18' -config['tags'] = ['review_7'] +config['name'] = 'curv_resnet_1' +config['notes'] = '' +config['tags'] = ['iccv', 'resnet'] config['cond_ae']['upscale_output'] = False config['pretrained_model'] = '' diff --git a/configs/iccv/config_ae_static.py b/configs/iccv/config_ae_static.py index 5890114..13a5c64 100644 --- a/configs/iccv/config_ae_static.py +++ b/configs/iccv/config_ae_static.py @@ -23,17 +23,17 @@ config['cond_ae']['is_label'] = False config['cond_ae']['upscale_output'] = False config['cond_ae']['latent_dim'] = config['latent_dim'] config['cond_ae']['kernel_size'] = 7 -config['cond_ae']['bottleneck_size'] = 8*4 # 8*4 -config['cond_ae']['layer_params'] = [3, 4, 23, 4, 3] -# config['cond_ae']['layer_params'] = [3, 4, 23, 3, 3] -config['cond_ae']['activation'] = 'linear' +config['cond_ae']['bottleneck_size'] = 8*2 # 8*4 +# config['cond_ae']['layer_params'] = [3, 4, 23, 4, 3] +config['cond_ae']['layer_params'] = [3, 7, 15, 7, 3] +config['cond_ae']['activation'] = 'tanh' -config['gpu'] = '5' -config['lr'] = 0.0001 +config['gpu'] = '1' +config['lr'] = 3e-4 config['name'] = 'static_resnet_1' -config['notes'] = 'resnet101 based, one more block' -config['tags'] = ['iccv', 'resnet'] +config['notes'] = '' +config['tags'] = ['iccv', 'resnet', 'static'] diff --git a/custom_utils/custom_utils.py b/custom_utils/custom_utils.py index db0fcf2..9e76a2c 100644 --- a/custom_utils/custom_utils.py +++ b/custom_utils/custom_utils.py @@ -432,4 +432,28 @@ def group_inputs_for_levels(batch, config): uv_orig = batch['uv_posed' + res_modifier] levels_input.append([uv_orig, batch['uv_static' + res_modifier], batch['uv_body_posed' + res_modifier], batch['uv_curvature' + res_modifier], batch['uv_mask' + res_modifier], batch['uv_mask']]) conds = [batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature']] - return levels_input, conds \ No newline at end of file + return levels_input, conds + + + + + +class Accumulator(): + def __init__(self): + self.dict = dict() + + def update(self, dict): + for key in dict.keys(): + if key in self.dict.keys(): + self.dict[key].append(dict[key]) + else: + self.dict[key] = [dict[key]] + + def get(self): + result = dict() + for key in self.dict.keys(): + result[key] = np.mean(self.dict[key]) + return result + + def reset(self): + self.dict = dict() \ No newline at end of file diff --git a/models/base.py b/models/base.py index 1b767c6..d077ecb 100644 --- a/models/base.py +++ b/models/base.py @@ -166,7 +166,7 @@ class AEEncoderNew(tf.keras.Model): def __init__(self, layer_params, features, activation='tanh'): super(AEEncoderNew, self).__init__() - self.encoder = resnet_18(layer_params, features) + self.encoder = resnet_101(layer_params, features) self.act = tf.keras.layers.Activation(activation) @@ -193,9 +193,9 @@ class DecoderNew(tf.keras.Model): self.is_label = is_label if not is_label: - self.decoder = resnet_18(layer_params, features, reversed=True, is_cvae=is_cvae) + self.decoder = resnet_101(layer_params, features, reversed=True, is_cvae=is_cvae) else: - self.decoder = resnet_18(layer_params, features, reversed=True, labels=6, is_cvae=is_cvae) + self.decoder = resnet_101(layer_params, features, reversed=True, labels=6, is_cvae=is_cvae) def call(self, inputs, training): diff --git a/models/resnet.py b/models/resnet.py index b32e3a4..ef72412 100644 --- a/models/resnet.py +++ b/models/resnet.py @@ -133,7 +133,7 @@ class ResNetTypeI(tf.keras.Model): self.labels = labels if not reversed: - self.conv1 = tf.keras.layers.Conv2D(filters=64, + self.conv1 = tf.keras.layers.Conv2D(filters=32, kernel_size=(7, 7), strides=2, padding="same") @@ -149,19 +149,19 @@ class ResNetTypeI(tf.keras.Model): self.bn1 = tf.keras.layers.BatchNormalization() self.pool1 = tf.keras.layers.UpSampling2D() - self.layer1 = make_basic_block_layer(filter_num=64, + self.layer1 = make_basic_block_layer(filter_num=32, blocks=layer_params[0], stride=1 if not reversed else 2, upsample=reversed) - self.layer2 = make_basic_block_layer(filter_num=64, + self.layer2 = make_basic_block_layer(filter_num=32, blocks=layer_params[1], stride=2, upsample=reversed) - self.layer3 = make_basic_block_layer(filter_num=64, + self.layer3 = make_basic_block_layer(filter_num=32, blocks=layer_params[2], stride=2, upsample=reversed) - self.layer4 = make_basic_block_layer(filter_num=128, + self.layer4 = make_basic_block_layer(filter_num=64, blocks=layer_params[3], stride=2, upsample=reversed) @@ -207,7 +207,7 @@ class ResNetTypeII(tf.keras.Model): self.labels = labels if not reversed: - self.conv1 = tf.keras.layers.Conv2D(filters=64, + self.conv1 = tf.keras.layers.Conv2D(filters=32, kernel_size=(7, 7), strides=2, padding="same") @@ -223,19 +223,19 @@ class ResNetTypeII(tf.keras.Model): self.bn1 = tf.keras.layers.BatchNormalization() self.pool1 = tf.keras.layers.UpSampling2D() - self.layer1 = make_bottleneck_layer(filter_num=64, + self.layer1 = make_bottleneck_layer(filter_num=32, blocks=layer_params[0], stride=1 if not reversed else 2, upsample=reversed) - self.layer2 = make_bottleneck_layer(filter_num=128, + self.layer2 = make_bottleneck_layer(filter_num=64, blocks=layer_params[1], stride=2, upsample=reversed) - self.layer3 = make_bottleneck_layer(filter_num=256, + self.layer3 = make_bottleneck_layer(filter_num=64, blocks=layer_params[2], stride=2, upsample=reversed) - self.layer4 = make_bottleneck_layer(filter_num=512, + self.layer4 = make_bottleneck_layer(filter_num=128, blocks=layer_params[3], stride=2, upsample=reversed)