Kurz vorm durchdrehen

lib/utils/logging.py

@@ -5,10 +5,12 @@ from pytorch_lightning.loggers.neptune import NeptuneLogger
 from pytorch_lightning.loggers.test_tube import TestTubeLogger
 
 from lib.utils.config import Config
-
+import numpy as np
 
 class Logger(LightningLoggerBase):
 
+    media_dir = 'media'
+
     @property
     def experiment(self):
         if self.debug:
@@ -84,7 +86,9 @@ class Logger(LightningLoggerBase):
 
     def log_image(self, name, image, **kwargs):
         self.neptunelogger.log_image(name, image, **kwargs)
-        image.savefig(self.log_dir / name)
+        step = kwargs.get('step', None)
+        name = f'{step}_{name}' if step is not None else name
+        image.savefig(self.log_dir / self.media_dir / name)
 
     def save(self):
         self.testtubelogger.save()