Implementation:AUTOMATIC1111 Stable diffusion webui LearnRateScheduler for hypernetwork

Knowledge Sources	stable-diffusion-webui
Domains	Deep Learning, Optimization, Training Configuration
Last Updated	2026-02-08 00:00 GMT

Overview

Concrete implementation of learning rate scheduling and optimizer configuration for hypernetwork training in Stable Diffusion, provided by the AUTOMATIC1111 stable-diffusion-webui repository. LearnRateScheduler parses a step-based learning rate schedule string and applies rate changes during training, while optimizer setup in the hypernetwork training function configures AdamW (or any PyTorch optimizer) with the hypernetwork's trainable parameters.

Description

The LearnRateScheduler class parses a comma-separated learning rate schedule string (e.g., "0.005:100, 0.0005:1000") into discrete phases. At each training step, calling apply(optimizer, step_number) checks whether a phase transition should occur and updates the optimizer's learning rate accordingly.

In the hypernetwork training context, the scheduler is used alongside the optimizer setup at hypernetwork.py:L549-559, which selects the optimizer class from torch.optim based on the hypernetwork's stored optimizer_name. The default optimizer is AdamW, and the optimizer state can be restored from a saved checkpoint for training resumption.

Gradient clipping is also configured in the training function, using a second LearnRateScheduler instance to schedule the clipping threshold value.

Usage

Import and use LearnRateScheduler when implementing the training loop for hypernetworks. It is instantiated at the beginning of train_hypernetwork() and called on every training step via its apply() method.

Code Reference

Source Location

• Repository: stable-diffusion-webui
• File: modules/textual_inversion/learn_schedule.py (LearnRateScheduler)
• Lines: L50-81
• File: modules/hypernetworks/hypernetwork.py (optimizer setup)
• Lines: L549-559

Signature

class LearnRateScheduler:
    def __init__(self, learn_rate, max_steps, cur_step=0, verbose=True):

    def step(self, step_number):
        """Returns True if the learning rate changed, False otherwise."""

    def apply(self, optimizer, step_number):
        """Calls step() and updates optimizer param_groups if rate changed."""

Optimizer setup in train_hypernetwork():

# hypernetwork.py L549-559
weights = hypernetwork.weights()
hypernetwork.train()

if hypernetwork.optimizer_name in optimizer_dict:
    optimizer = optimizer_dict[hypernetwork.optimizer_name](params=weights, lr=scheduler.learn_rate)
    optimizer_name = hypernetwork.optimizer_name
else:
    print(f"Optimizer type {hypernetwork.optimizer_name} is not defined!")
    optimizer = torch.optim.AdamW(params=weights, lr=scheduler.learn_rate)
    optimizer_name = 'AdamW'

Import

from modules.textual_inversion.learn_schedule import LearnRateScheduler

I/O Contract

Inputs (LearnRateScheduler.__init__)

Name	Type	Required	Description
learn_rate	str	Yes	Learning rate schedule string, e.g., "0.005:100, 0.0005:1000" or a single value like "0.001"
max_steps	int	Yes	Total number of training steps; phases are clamped to this value
cur_step	int	No	Current step to resume from; phases ending before this step are skipped (default: 0)
verbose	bool	No	Whether to print learning rate transitions to console (default: True)

Inputs (LearnRateScheduler.apply)

Name	Type	Required	Description
optimizer	torch.optim.Optimizer	Yes	The optimizer whose param_groups will have their lr updated
step_number	int	Yes	The current training step number

Outputs

Name	Type	Description
self.learn_rate	float	The current active learning rate
self.finished	bool	True when all schedule phases have been exhausted
self.end_step	int	The step number at which the current phase ends

Usage Examples

Basic Learning Rate Scheduling

from modules.textual_inversion.learn_schedule import LearnRateScheduler
import torch

# Multi-phase schedule: 0.005 until step 100, then 0.0005 until step 1000
scheduler = LearnRateScheduler("0.005:100, 0.0005:1000", max_steps=1000, cur_step=0)

# Create optimizer with hypernetwork weights
optimizer = torch.optim.AdamW(params=hypernetwork.weights(), lr=scheduler.learn_rate)

for step in range(1000):
    if scheduler.finished:
        break
    scheduler.apply(optimizer, step)
    # optimizer param_groups now have the correct lr for this step
    # ... training logic ...

Gradient Clipping with Scheduled Threshold

from modules.textual_inversion.learn_schedule import LearnRateScheduler
import torch

# Main learning rate schedule
scheduler = LearnRateScheduler("0.005:500, 0.0005:5000", max_steps=5000)

# Gradient clip schedule (reusing LearnRateScheduler for the threshold value)
clip_grad_sched = LearnRateScheduler("1.0:5000", max_steps=5000, verbose=False)

clip_grad = torch.nn.utils.clip_grad_norm_

for step in range(5000):
    scheduler.apply(optimizer, step)
    clip_grad_sched.step(step)
    # ... forward + backward ...
    clip_grad(hypernetwork.weights(), clip_grad_sched.learn_rate)
    optimizer.step()

Resuming Training from Checkpoint

from modules.textual_inversion.learn_schedule import LearnRateScheduler

# Resume from step 500
initial_step = hypernetwork.step  # e.g., 500
scheduler = LearnRateScheduler("0.005:100, 0.0005:1000", max_steps=1000, cur_step=initial_step)
# Phase "0.005:100" is skipped since 100 < 500
# Starts at phase "0.0005:1000"

Related Pages

Implements Principle

• Principle:AUTOMATIC1111_Stable_diffusion_webui_Hypernetwork_training_configuration