update swanlab log

This commit is contained in:
Artiprocher
2025-02-27 19:21:51 +08:00
parent 1419bec53d
commit b5c1d33e58
2 changed files with 14 additions and 46 deletions

View File

@@ -256,29 +256,11 @@ def add_general_parsers(parser):
action="store_true",
help="Whether to use SwanLab logger.",
)
parser.add_argument(
"--swanlab_project",
type=str,
default="diffsynth_studio",
help="SwanLab project name.",
)
parser.add_argument(
"--swanlab_name",
type=str,
default="diffsynth_studio_train",
help="SwanLab experimentname.",
)
parser.add_argument(
"--swanlab_mode",
default=None,
help="SwanLab mode (cloud or local).",
)
parser.add_argument(
"--swanlab_logdir",
type=str,
default=None,
help="SwanLab local log directory.",
)
return parser
@@ -301,18 +283,20 @@ def launch_training_task(model, args):
)
# set swanlab logger
swanlab_logger = None
if args.use_swanlab:
from swanlab.integration.pytorch_lightning import SwanLabLogger
swanlab_config = {"UPPERFRAMEWORK": "DiffSynth-Studio"}
swanlab_config.update(vars(args))
swanlab_logger = SwanLabLogger(
project=args.swanlab_project,
name=args.swanlab_name,
project="diffsynth_studio",
name="diffsynth_studio",
config=swanlab_config,
mode=args.swanlab_mode,
logdir=args.swanlab_logdir,
logdir=args.output_path,
)
logger = [swanlab_config]
else:
logger = []
# train
trainer = pl.Trainer(
@@ -324,7 +308,7 @@ def launch_training_task(model, args):
default_root_dir=args.output_path,
accumulate_grad_batches=args.accumulate_grad_batches,
callbacks=[pl.pytorch.callbacks.ModelCheckpoint(save_top_k=-1)],
logger=[swanlab_logger],
logger=logger,
)
trainer.fit(model=model, train_dataloaders=train_loader)