kohya-ss / kohya-ss/sd-scripts

[Feature] Supporting individual learning rates for multiple TEs

Open
#935 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

(Edit: upon further testing I think this implementation is currently broken. I will be trying to make the modifications necessary to try and get this working)
I just got this implemented in my own installation, and 0 changes needed to be made to sdxl_train_network.py. However, I am using the bmaltais/kohya_ss GUI, and I had to make a few changes to lora_gui.py as well to get it working. Overall this is a pretty easy change to make and doesn't seem to break any current functionality (single learning rate still works, and all parsing in that case is the same, text_encoder_lr is now just a list behind the scenes).
![image](https://github.com/kohya-ss/sd-scripts/assets/31860133/459a2c4c-9c92-4c46-b600-07ae44b3c717)
I dont know if I want to make a pull request since its a pretty small change and I only have the GUI version of this repo installed locally, so here are the only changes required to get this to work:

In train_network.py:
```py
...
# 学習に必要なクラスを準備する
accelerator.print("prepare optimizer, data loader etc.")

# 後方互換性を確保するよ
try: #Normal usage
if len(args.text_encoder_lr) == 1:
trainable_params = network.prepare_optimizer_params(float(args.text_encoder_lr[0]), args.unet_lr, args.learning_rate)
else:
trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr, args.learning_rate)

except TypeError: #Deprecated usage (omitted learning_rate)
accelerator.print(
"Deprecated: use prepare_optimizer_params(text_encoder_lr, unet_lr, learning_rate) instead of prepare_optimizer_params(text_encoder_lr, unet_lr)"
)
if len(args.text_encoder_lr) == 1:
trainable_params = network.prepare_optimizer_params(float(args.text_encoder_lr[0]), args.unet_lr)
else:
trainable_params = network.prepare_optimizer_params(args.text_encoder_lr, args.unet_lr)

optimizer_name, optimizer_args, optimizer = train_util.get_optimizer(args, trainable_params)
...
#in setup_parser, change the type to a list:
parser.add_argument("--text_encoder_lr", type=list, default=None, help="learning rate for Text Encoder(s) / Text Encoder(s)の学習率")
...
```
In lora.py:
```py
# 二つのText Encoderに別々の学習率を設定できるようにするといいかも
def prepare_optimizer_params(self, text_encoder_lr, unet_lr, default_lr):
self.requires_grad_(True)
all_params = []

def enumerate_params(loras):
params = []
for lora in loras:
params.extend(lora.parameters())
return params

if self.text_encoder_loras:
param_data = {"params": enumerate_params(self.text_encoder_loras)}
#NEW CODE:
if len(text_encoder_lr) > 1: #multiple TE Lrs provided
#assert len(param_data["params"])==len(text_encoder_lr)
for i in range(0, len(text_encoder_lr)):
if text_encoder_lr[i] is not None:
param_data["lr"] = float(text_encoder_lr[i])
all_params.append(param_data)
#ORIGINAL PROCESSING:
else:
if text_encoder_lr is not None:
param_data["lr"] = float(text_encoder_lr[i])
all_params.append(param_data)
```
This implementation would require users of the script to update old configs; adding square brackets around the original float values (e.g. [1.0] instead of 1.0 ). text_encoder_lr would now be given as a list [lr_1, ..., lr_n], applying the lrs sequentially to the enumerated params. By making some small changes to the GUI, this is done implicitly and hidden from basic users if they only wish to input a single LR. An improved implementation of this would seamlessly add the ability to specify different learning rates for each model. Please consider adding something similar to this feature.

P.S.
I have a theory that when using adaptive optimizers like Prodigy, having the LR set to 1.0 for both of them results in an effective LR of 2.0, causing rapid burnup. I think being able to set the LR individually for each model (to something like 1/n strength each) opens up the doors to finally fixing the SDXL TE training problem. What do you think?

_Originally posted by @antis0007 in https://github.com/kohya-ss/sd-scripts/issues/930#issuecomment-1800691707_

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading the argument parsing and optimizer setup in train_network.py, then inspect prepare_optimizer_params in lora.py. Compare the existing behavior with sdxl_train_network.py and the GUI changes described in the issue. Done means separate learning rates work for multiple text encoders while the existing single-rate behavior remains functional.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
28/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.