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)

View File

@@ -400,29 +400,11 @@ def parse_args():
action="store_true",
help="Whether to use SwanLab logger.",
)
parser.add_argument(
"--swanlab_project",
type=str,
default="wan_t2v",
help="SwanLab project name.",
)
parser.add_argument(
"--swanlab_name",
type=str,
default="wan_t2v_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.",
)
args = parser.parse_args()
return args
@@ -493,18 +475,20 @@ def train(args):
init_lora_weights=args.init_lora_weights,
use_gradient_checkpointing=args.use_gradient_checkpointing
)
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="wan",
name="wan",
config=swanlab_config,
mode=args.swanlab_mode,
logdir=args.swanlab_logdir,
logdir=args.output_path,
)
logger = [swanlab_logger]
else:
logger = []
trainer = pl.Trainer(
max_epochs=args.max_epochs,
accelerator="gpu",
@@ -513,7 +497,7 @@ def train(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, dataloader)