github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Adapter work. Bug fixes. Auto adjust LR when resuming optimizer.

Browse Source
This commit is contained in:
Jaret Burkett
2024-03-17 10:21:47 -06:00
parent 72de68d8aa
commit 016687bda1
8 changed files with 84 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
jobs/process/BaseSDTrainProcess.py
View File

@@ -1352,6 +1352,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
# try to load
# previous param groups
# previous_params = copy.deepcopy(optimizer.param_groups)
previous_lrs = []
for group in optimizer.param_groups:
previous_lrs.append(group['lr'])
try:
print(f"Loading optimizer state from {optimizer_state_file_path}")
optimizer_state_dict = torch.load(optimizer_state_file_path)
@@ -1360,6 +1364,13 @@ class BaseSDTrainProcess(BaseTrainProcess):
print(f"Failed to load optimizer state from {optimizer_state_file_path}")
print(e)
# update the optimizer LR from the params
print(f"Updating optimizer LR from params")
if len(previous_lrs) > 0:
for i, group in enumerate(optimizer.param_groups):
group['lr'] = previous_lrs[i]
group['initial_lr'] = previous_lrs[i]
# Update the learning rates if they changed
# optimizer.param_groups = previous_params