Allow logging models from GenericLogger

classifier.py

@@ -249,7 +249,9 @@ def train():
         images = images.to(device)
         pred = torch.max(ema.ema(images), 1)[1]
         file = imshow(denormalize(images), labels, pred, names, verbose=True, f=save_dir / 'test_images.jpg')
+        meta = {"epochs": epochs, "accuracy": best_fitness, "date": datetime.now().isoformat()}
         logger.log_images(file, name='Test Examples (true-predicted)', epoch=epoch)
+        logger.log_model(best, epochs, metadata=meta)
 
 
 @torch.no_grad()

utils/loggers/__init__.py

@@ -239,6 +239,13 @@ def log_graph(self, model, imgsz=(640, 640)):
         if self.tb:
             log_tensorboard_graph(self.tb, model, imgsz)
 
+    def log_model(self, model_path, epoch=0, metadata={}):
+        # Log model to all loggers
+        if self.wandb:
+            art = wandb.Artifact(name=f"run_{wandb.run.id}_model", type="model", metadata=metadata)
+            art.add_file(str(model_path))
+            wandb.log_artifact(art)
+
 
 def log_tensorboard_graph(tb, model, imgsz=(640, 640)):
     # Log model graph to TensorBoard