apply for body and curv

This commit is contained in:
2023-01-25 19:13:18 +00:00
parent 8514032837
commit 655c70413b
7 changed files with 125 additions and 17 deletions
+7 -4
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.config_ae_body_post import config
from configs.review.config_ae_body import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -34,7 +34,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
def init_wandb():
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp')
run = wandb.init(project='ae', name=config['name'], notes=config['notes'], dir='/home/workspace/local-data/temp', tags=config['tags'])
artifact = wandb.Artifact('code', type='code')
artifact.add_file('auto_encoder_body.py')
artifact.add_file('losses.py')
@@ -42,7 +42,7 @@ 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.')]
config_file_names = [key for key in sys.modules.keys() if key.startswith('configs.review.')]
for config_file_name in config_file_names:
artifact.add_file('/'.join(config_file_name.split('.')) + '.py')
run.log_artifact(artifact)
@@ -137,7 +137,10 @@ def train():
if step == 0:
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
wandb.log(to_log, step=step)
if step % 1000:
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)