From 40c19f3939426aa388f0955229efdb75442b0f83 Mon Sep 17 00:00:00 2001 From: Jiao-Long Cao Date: Tue, 7 Jan 2025 16:43:34 +0800 Subject: [PATCH] Add wandb project name argument and allow change wandb run name --- train.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/train.py b/train.py index a9e63f45..182b45b8 100755 --- a/train.py +++ b/train.py @@ -388,6 +388,8 @@ group.add_argument('--use-multi-epochs-loader', action='store_true', default=Fal help='use the multi-epochs-loader to save time at the beginning of every epoch') group.add_argument('--log-wandb', action='store_true', default=False, help='log training and validation metrics to wandb') +group.add_argument('--wandb-project', default=None, type=str, + help='wandb project name') group.add_argument('--wandb-tags', default=[], type=str, nargs='+', help='wandb tags') group.add_argument('--wandb-resume-id', default='', type=str, metavar='ID', @@ -823,9 +825,14 @@ def main(): if utils.is_primary(args) and args.log_wandb: if has_wandb: assert not args.wandb_resume_id or args.resume - wandb.init(project=args.experiment, config=args, tags=args.wandb_tags, - resume='must' if args.wandb_resume_id else None, - id=args.wandb_resume_id if args.wandb_resume_id else None) + wandb.init( + project=args.wandb_project, + name=args.experiment, + config=args, + tags=args.wandb_tags, + resume="must" if args.wandb_resume_id else None, + id=args.wandb_resume_id if args.wandb_resume_id else None, + ) else: _logger.warning( "You've requested to log metrics to wandb but package not found. "