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

update

Browse Source
This commit is contained in:
luni
2022-09-13 09:23:01 +00:00
parent 977e49900d
commit 440444cd82
19 changed files with 1630 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
auto_encoder_body.py
+10 -10
View File
@@ -25,9 +25,9 @@ from dataset import DatasetFactory, StaticDataset
from custom_utils.custom_utils import batch_to_dict
from custom_utils.smpl import SMPL
name='ae_body_7'
name='ae_body_7_2'
def init_wandb():
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
run = wandb.init(project='vae', name=name,
notes='convnext + convnext, add dense with tanh at the end of reverse convnext, larger features. add lr scheduler. fft3d instead of rfft3d. \
fft2d per channel and separate for top and bottom. fft_loss, all data. added convnext block as bottleneck, same as before, added extra metric. \
@@ -35,7 +35,7 @@ def init_wandb():
two more downscaling than up, so decoder output is resized wto final resolution\
larger batch_size.\
flattened bottleneck with 128 latent_dim\ remove global orientation\
fixed gender issue with deoriented bodies, again\
fixed gender issue with deoriented bodies, again, aaand again\
trained on 100 garments')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
@@ -377,14 +377,14 @@ def train():
# data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, signature_dynamic))
valid_loss = {}#valid_step(data_valid)
if step % 1000 == 0:
if step % 2500 == 0:
train_figs = plot_everything(batch, 'train')
valid_figs = {}#plot_everything(data_valid, 'valid')
wandb.log({ **train_loss, **valid_loss, **train_figs, **valid_figs }, step=step)
else:
wandb.log({ **train_loss, **valid_loss }, step=step)
if step % 2500 == 0:
if step % 10000 == 0:
manager.save(checkpoint_number=step)
@@ -410,8 +410,8 @@ if __name__ == '__main__':
# 'uv_posed': (3, 512, 256,),
'uv_body_posed': (3, 512, 256,),
'mesh_body_posed': (14475, 3,),
'mesh_posed': (14475, 3,),
'mesh_unposed': (14475, 3,),
# 'mesh_posed': (14475, 3,),
# 'mesh_unposed': (14475, 3,),
# 'uv_unposed': (3, 512, 256,)
}
signature_static = { #'uv_static': (3, 512, 256,),
@@ -421,8 +421,8 @@ if __name__ == '__main__':
# 'uv_shaped_body': (3, 512, 256,),
}
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static)
dataset_train = factory.get_dataset(folder='/home/workspace/local-data/train_100', batch_size=32, signature=signature_dynamic, use_filled=False)
static_dataset_train = StaticDataset(folder='/home/workspace/local-data/train_100', signature=signature_static, use_filled=False)
dataset_valid = factory.get_dataset(folder='/home/workspace/local-data/jumpsuit_20/valid', batch_size=32, signature=signature_dynamic)
static_dataset_valid = StaticDataset(folder='/home/workspace/local-data/jumpsuit_20/valid', signature=signature_static)
@@ -432,6 +432,6 @@ if __name__ == '__main__':
folder = 'models/ae_logs/'
ckpt = tf.train.Checkpoint(model=model)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=5)
manager = tf.train.CheckpointManager(ckpt, 'checkpoints/' + name, max_to_keep=3)
train()