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

Added proper grad accumulation

Browse Source
This commit is contained in:
Jaret Burkett
2024-09-03 07:24:18 -06:00
parent e5fadddd45
commit 121a760c19
3 changed files with 78 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
toolkit/config_modules.py
View File

@@ -236,6 +236,7 @@ class TrainConfig:
self.min_denoising_steps: int = kwargs.get('min_denoising_steps', 0)
self.max_denoising_steps: int = kwargs.get('max_denoising_steps', 1000)
self.batch_size: int = kwargs.get('batch_size', 1)
self.orig_batch_size: int = self.batch_size
self.dtype: str = kwargs.get('dtype', 'fp32')
self.xformers = kwargs.get('xformers', False)
self.sdp = kwargs.get('sdp', False)
@@ -284,8 +285,16 @@ class TrainConfig:
# set to -1 to accumulate gradients for entire epoch
# warning, only do this with a small dataset or you will run out of memory
# This is legacy but left in for backwards compatibility
self.gradient_accumulation_steps = kwargs.get('gradient_accumulation_steps', 1)
# this will do proper gradient accumulation where you will not see a step until the end of the accumulation
# the method above will show a step every accumulation
self.gradient_accumulation = kwargs.get('gradient_accumulation', 1)
if self.gradient_accumulation > 1:
if self.gradient_accumulation_steps != 1:
raise ValueError("gradient_accumulation and gradient_accumulation_steps are mutually exclusive")
# short long captions will double your batch size. This only works when a dataset is
# prepared with a json caption file that has both short and long captions in it. It will
# Double up every image and run it through with both short and long captions. The idea