mirror of
https://github.com/modelscope/DiffSynth-Studio.git
synced 2026-03-18 22:08:13 +00:00
update swanlab log
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user