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

fixed bug the kept learning rates the same

Browse Source
This commit is contained in:
Jaret Burkett
2024-03-06 09:23:32 -07:00
parent b01e8d889a
commit f1cb87fe9e
2 changed files with 18 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extensions_built_in/sd_trainer/SDTrainer.py
View File

@@ -1363,7 +1363,13 @@ class SDTrainer(BaseSDTrainProcess):
# flush() # flush()
if not self.is_grad_accumulation_step: if not self.is_grad_accumulation_step:
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm) # torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
# fix this for multi params
if isinstance(self.params[0], dict):
for i in range(len(self.params)):
torch.nn.utils.clip_grad_norm_(self.params[i]['params'], self.train_config.max_grad_norm)
else:
torch.nn.utils.clip_grad_norm_(self.params, self.train_config.max_grad_norm)
# only step if we are not accumulating # only step if we are not accumulating
with self.timer('optimizer_step'): with self.timer('optimizer_step'):
# apply gradients # apply gradients

18
jobs/process/BaseSDTrainProcess.py
View File

@@ -1327,13 +1327,14 @@ class BaseSDTrainProcess(BaseTrainProcess):
flush() flush()
### HOOK ### ### HOOK ###
params = self.hook_add_extra_train_params(params) params = self.hook_add_extra_train_params(params)
self.params = [] self.params = params
# self.params = []
for param in params: # for param in params:
if isinstance(param, dict): # if isinstance(param, dict):
self.params += param['params'] # self.params += param['params']
else: # else:
self.params.append(param) # self.params.append(param)
if self.train_config.start_step is not None: if self.train_config.start_step is not None:
self.step_num = self.train_config.start_step self.step_num = self.train_config.start_step
@@ -1431,6 +1432,9 @@ class BaseSDTrainProcess(BaseTrainProcess):
flush() flush()
# self.step_num = 0 # self.step_num = 0
print(f"Compiling Model")
torch.compile(self.sd.unet, dynamic=True)
################################################################### ###################################################################
# TRAIN LOOP # TRAIN LOOP
################################################################### ###################################################################
@@ -1438,6 +1442,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
start_step_num = self.step_num start_step_num = self.step_num
did_first_flush = False did_first_flush = False
for step in range(start_step_num, self.train_config.steps): for step in range(start_step_num, self.train_config.steps):
self.timer.start('train_loop')
if self.train_config.do_random_cfg: if self.train_config.do_random_cfg:
self.train_config.do_cfg = True self.train_config.do_cfg = True
self.train_config.cfg_scale = value_map(random.random(), 0, 1, 1.0, self.train_config.max_cfg_scale) self.train_config.cfg_scale = value_map(random.random(), 0, 1, 1.0, self.train_config.max_cfg_scale)
@@ -1506,7 +1511,6 @@ class BaseSDTrainProcess(BaseTrainProcess):
# flush() # flush()
### HOOK ### ### HOOK ###
self.timer.start('train_loop')
loss_dict = self.hook_train_loop(batch) loss_dict = self.hook_train_loop(batch)
self.timer.stop('train_loop') self.timer.stop('train_loop')
if not did_first_flush: if not did_first_flush: