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

working auto encoders

Browse Source
This commit is contained in:
luni
2022-09-25 22:12:36 +00:00
parent 26916c00e7
commit 040b29e898
9 changed files with 41 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+10 -7
View File
@@ -31,7 +31,7 @@ def init_wandb():
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
artifact.add_file('configs/config.py')
artifact.add_file('configs/config_base.py')
artifact.add_file('configs/config_ae_body.py')
artifact.add_file('models/base.py')
artifact.add_file('models/cond_encoder.py')
@@ -45,7 +45,7 @@ uv2mesh = UV2Mesh(res=256)
@tf.function
def train_step(batch):
with tf.GradientTape() as tape:
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=True)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
@@ -87,7 +87,7 @@ def valid_step(batch):
def plot_everything(batch: dict, type: str):
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body])
pred = model([batch['uv_body_posed' + resolution_modifier], uv_mask_body], training=False)
mesh_pred = uv2mesh(pred)
mesh_orig = batch['mesh_body_posed']
@@ -133,7 +133,8 @@ def train():
wandb.log({ **train_loss, **valid_loss }, step=step)
if step % 10000 == 0:
manager.save(checkpoint_number=step)
manager_enc.save(checkpoint_number=step)
manager_dec.save(checkpoint_number=step)
print(step)
@@ -141,7 +142,7 @@ def train():
if __name__ == '__main__':
# init_wandb()
init_wandb()
allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
@@ -165,7 +166,9 @@ if __name__ == '__main__':
model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + config['name'], max_to_keep=3)
ckpt_enc = tf.train.Checkpoint(model=model.encoder)
manager_enc = tf.train.CheckpointManager(ckpt_enc, 'checkpoints/' + config['name'] + '_enc', max_to_keep=3)
ckpt_dec = tf.train.Checkpoint(model=model.decoder)
manager_dec = tf.train.CheckpointManager(ckpt_dec, 'checkpoints/' + config['name'] + '_dec', max_to_keep=3)
train()