ae updates

This commit is contained in:
2022-10-19 18:42:59 +00:00
parent 8eaf090e4c
commit 1f85598db5
3 changed files with 27 additions and 29 deletions
+6 -6
View File
@@ -35,7 +35,7 @@ from models.cond_encoder import AutoEncoder
testing = False
def init_wandb():
if not testing:
run = wandb.init(project='ae', name=config['name'], notes=config['notes'])
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp/')
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_curv.py')
artifact.add_file('losses.py')
@@ -47,7 +47,7 @@ def init_wandb():
artifact.add_file('convnext/models/convnext_tf.py')
run.log_artifact(artifact)
uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'])
uv2mesh = UV2MeshLabel(res=256, is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape'])
g_mean_iou = tf.keras.metrics.MeanIoU(num_classes=7)
@@ -202,10 +202,10 @@ if __name__ == '__main__':
factory = DatasetFactory()
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True)
dataset_train = factory.get_dataset(folder=config['train_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_train = StaticDataset(folder=config['train_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
dataset_valid = factory.get_dataset(folder=config['valid_dir'], batch_size=config['batch_size'], signature=config['signature_dynamic'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=config['dataset_in_memory'], is_cape=config['is_cape'])
static_dataset_valid = StaticDataset(folder=config['valid_dir'], signature=config['signature_static'], use_filled=config['use_filled'], resolutions=[config['resolution']], in_memory=True, is_cape=config['is_cape'])
model = AutoEncoder(config['cond_ae'])
optimizer = tf.optimizers.Adam(learning_rate=config['lr'])