ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-05-01 03:31:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed issues with vgg19 preprocessing. Added yaml support on config file

Browse Source
This commit is contained in:
Jaret Burkett
2023-07-19 06:12:57 -06:00
parent 8d8229dfc0
commit 6ada328d8d
6 changed files with 46 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
jobs/TrainJob.py
View File

@@ -26,7 +26,7 @@ class TrainJob(BaseJob):
self.device = self.get_conf('device', 'cpu')
self.gradient_accumulation_steps = self.get_conf('gradient_accumulation_steps', 1)
self.mixed_precision = self.get_conf('mixed_precision', False) # fp16
self.logging_dir = self.get_conf('logging_dir', None)
self.log_dir = self.get_conf('log_dir', None)
self.writer = None
self.setup_tensorboard()
@@ -43,9 +43,9 @@ class TrainJob(BaseJob):
process.run()
def setup_tensorboard(self):
if self.logging_dir:
if self.log_dir:
from torch.utils.tensorboard import SummaryWriter
self.writer = SummaryWriter(
log_dir=self.logging_dir,
log_dir=self.log_dir,
filename_suffix=f"_{self.name}"
)