TensorBoard `on_train_end()` speed improvements (#4605)
Browse files
utils/loggers/__init__.py
CHANGED
@@ -131,10 +131,9 @@ class Loggers():
|
|
131 |
files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
|
132 |
|
133 |
if self.tb:
|
134 |
-
|
135 |
-
import numpy as np
|
136 |
for f in files:
|
137 |
-
self.tb.add_image(f.stem,
|
138 |
|
139 |
if self.wandb:
|
140 |
self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
|
|
|
131 |
files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter
|
132 |
|
133 |
if self.tb:
|
134 |
+
import cv2
|
|
|
135 |
for f in files:
|
136 |
+
self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC')
|
137 |
|
138 |
if self.wandb:
|
139 |
self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]})
|