ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-01-26 16:39:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issue with new logger when ooming

Browse Source
This commit is contained in:
Jaret Burkett
2025-12-25 16:57:34 +00:00
parent 8edf1e44c5
commit a767b82b60
2 changed files with 17 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
jobs/process/BaseSDTrainProcess.py
View File

@@ -2212,6 +2212,7 @@ class BaseSDTrainProcess(BaseTrainProcess):
with torch.no_grad():
# torch.cuda.empty_cache()
# if optimizer has get_lrs method, then use it
learning_rate = 0.0
if not did_oom and loss_dict is not None:
if hasattr(optimizer, 'get_avg_learning_rate'):
learning_rate = optimizer.get_avg_learning_rate()
@@ -2282,9 +2283,10 @@ class BaseSDTrainProcess(BaseTrainProcess):
# log to tensorboard
if self.accelerator.is_main_process:
if self.writer is not None:
for key, value in loss_dict.items():
self.writer.add_scalar(f"{key}", value, self.step_num)
self.writer.add_scalar(f"lr", learning_rate, self.step_num)
if loss_dict is not None:
for key, value in loss_dict.items():
self.writer.add_scalar(f"{key}", value, self.step_num)
self.writer.add_scalar(f"lr", learning_rate, self.step_num)
if self.progress_bar is not None:
self.progress_bar.unpause()
@@ -2293,10 +2295,11 @@ class BaseSDTrainProcess(BaseTrainProcess):
self.logger.log({
'learning_rate': learning_rate,
})
for key, value in loss_dict.items():
self.logger.log({
f'loss/{key}': value,
})
if loss_dict is not None:
for key, value in loss_dict.items():
self.logger.log({
f'loss/{key}': value,
})
elif self.logging_config.log_every is None:
if self.accelerator.is_main_process:
# log every step