Adding the piecewise_constant LR Scheduler

Please, I would like to ask permission to modify your code a little bit.

I was digging through Kohya's codes looking for an option that would give me more specific control over the LR, like, to change it from time to time or end at a value other than 0.

I ended up discovering that Kohya has the 'piecewise_constant' scheduler implemented, which fits perfectly for what I wanted.

To use it in your GUI, I just had to add the name to the LR Scheduler dropdown, and Kohya's scripts would take care of the rest perfectly.

I took the usage explanation directly from the diffusers code:

Create a schedule with a constant learning rate, using the learning rate set in optimizer.

Args:
    optimizer ([`~torch.optim.Optimizer`]):
        The optimizer for which to schedule the learning rate.
    step_rules (`string`):
        The rules for the learning rate. ex: step_rules='1:10,0.1:20,0.01:30,0.005' it means that the learning rate
        if multiple 1 for the first 10 steps, mutiple 0.1 for the next 20 steps, multiple 0.01 for the next 30
        steps and multiple 0.005 for the other steps.
    last_epoch (`int`, *optional*, defaults to -1):
        The index of the last epoch when resuming training.

Return:
    `torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.

The steps rules should go in 'LR scheduler extra arguments' on GUI.

I stumbled a bit trying to understand exactly how the rules worked, but using the code example, it's basically like this:

step_rules='1:10,0.1:20,0.01:30,0.005'

Until step 10, the LR will be 1x the defined LR. Then until step 20, it will be 0.1x the defined LR. Then until step 30, it will be 0.01x the defined LR, and then for the rest of the training, it will be 0.005x the defined LR. In the last multiplication, the multiplier is defined without the step, to indicate that it's for the rest of the training.

_PS.: I also had problems because the arg in code is step_rules, but in the comment block explanation it was showing rulesstep. I also had to change the double quotes to single quotes because it was giving an error at training start. 😅

bmaltais / kohya_ss

Adding the piecewise_constant LR Scheduler #2565