增加runid

This commit is contained in:
somunslotus 2025-12-24 17:02:05 +08:00
parent b5d11fdea4
commit bea5e4100a
1 changed files with 12 additions and 3 deletions

View File

@ -203,6 +203,15 @@ if __name__ == '__main__':
# 创建 Aim logger如果可用
if AIM_AVAILABLE:
# 从环境变量读取 Aim repo 地址,如果不存在则使用本地路径
aim_repo = os.environ.get('EXPERIMENT_REMOTE_REPO', os.path.join(model_output_dir, '.aim'))
print(f"Aim repo: {aim_repo}")
# 从环境变量读取 EXPERIMENT_RUN_ID
experiment_run_id = os.environ.get('EXPERIMENT_RUN_ID')
if experiment_run_id:
print(f"EXPERIMENT_RUN_ID: {experiment_run_id}")
# 准备超参数字典
hparams = {
'batch_size': args.batch_size,
@ -219,9 +228,9 @@ if __name__ == '__main__':
'every_n_epochs': args.every_n_epochs,
}
# 从环境变量读取 Aim repo 地址,如果不存在则使用本地路径
aim_repo = os.environ.get('EXPERIMENT_REMOTE_REPO', os.path.join(model_output_dir, '.aim'))
print(f"Aim repo: {aim_repo}")
# 如果存在 EXPERIMENT_RUN_ID添加到超参数中
if experiment_run_id:
hparams['id'] = experiment_run_id
aim_logger = AimLogger(
repo=aim_repo, # 从环境变量读取或使用本地路径