Allow logging models from GenericLogger (#8676)

* enhance

* revert

* allow training from scratch

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update --img argument from train.py 

single line

* fix image size from 640 to 128

* suport custom dataloader and augmentation

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* format

* Update dataloaders.py

* Single line return, single line comment, remove unused argument

* address PR comments

* fix spelling

* don't augment eval set

* use fstring

* update augmentations.py

* new maning convention for transforms

* reverse if statement, inline ops

* reverse if statement, inline ops

* updates

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* update dataloaders

* Remove additional if statement

* Remove is_train as redundant

* Cleanup

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Cleanup2

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update classifier.py

* Update augmentations.py

* fix: imshow clip warning

* update

* Revert ToTensorV2 removal

* Update classifier.py

* Update normalize values, revert uint8

* normalize image using cv2

* remove dedundant comment

* Update classifier.py

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* replace print with logger

* commit steps

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* support final model logging

* update

* update

* update

* update

* remove curses

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update classifier.py

* Update __init__.py

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

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