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
+5 -5
View File
@@ -1,6 +1,6 @@
import os
import sys
from configs.config_ae_curv_post import config
from configs.review.config_ae_curv import config
os.environ["CUDA_VISIBLE_DEVICES"] = config['gpu']
import tensorflow as tf
@@ -37,7 +37,7 @@ from models.cond_encoder import AutoEncoderNew as AutoEncoder
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/')
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_curv.py')
artifact.add_file('losses.py')
@@ -45,7 +45,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)
@@ -178,7 +178,7 @@ def train():
to_log = { **to_log, **{ 'model/summary': wandb.Html(get_summary_html(model)) } }
wandb.log(to_log, step=step)
print(step)
# print(step)
if step % 10000 == 0:
manager_enc.save(checkpoint_number=step)
@@ -206,7 +206,7 @@ def test():
if __name__ == '__main__':
init_wandb()
allow_memory_growth()
# allow_memory_growth()
parser = argparse.ArgumentParser(description='')
parser.add_argument('--name', type=str)
args = vars(parser.parse_args())