VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-08 22:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix to work lienar/cosine lr scheduler closes #1602 ref #1393

Browse Source
This commit is contained in:
Kohya S
2024-09-16 10:42:09 +09:00
parent be078bdaca
commit 96c677b459
1 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
library/train_util.py
View File

@@ -4707,6 +4707,15 @@ def get_scheduler_fix(args, optimizer: Optimizer, num_processes: int):
**lr_scheduler_kwargs, **lr_scheduler_kwargs,
) )
# these schedulers do not require `num_decay_steps`
if name == SchedulerType.LINEAR or name == SchedulerType.COSINE:
return schedule_func(
optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=num_training_steps,
**lr_scheduler_kwargs,
)
# All other schedulers require `num_decay_steps` # All other schedulers require `num_decay_steps`
if num_decay_steps is None: if num_decay_steps is None:
raise ValueError(f"{name} requires `num_decay_steps`, please provide that argument.") raise ValueError(f"{name} requires `num_decay_steps`, please provide that argument.")
@@ -5837,14 +5846,9 @@ def sample_image_inference(
wandb_tracker = accelerator.get_tracker("wandb") wandb_tracker = accelerator.get_tracker("wandb")
import wandb import wandb
# not to commit images to avoid inconsistency between training and logging steps # not to commit images to avoid inconsistency between training and logging steps
wandb_tracker.log( wandb_tracker.log({f"sample_{i}": wandb.Image(image, caption=prompt)}, commit=False) # positive prompt as a caption
{f"sample_{i}": wandb.Image(
image,
caption=prompt # positive prompt as a caption
)},
commit=False
)
# endregion # endregion