From 99f8b1b4b238888d68e8fa2e9b588c29780c59bf Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 18 Aug 2022 21:45:11 +0200
Subject: [PATCH] Add PyTorch Hub classification CI checks (#9027)

* Add PyTorch Hub classification CI checks

Add PyTorch Hub loading of official and custom trained classification models to CI checks.

May help resolve https://github.com/ultralytics/yolov5/issues/8790#issuecomment-1219840718

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

* Update hubconf.py

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
 .github/workflows/ci-testing.yml | 5 +++++
 hubconf.py                       | 6 +++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml
index aa797c44d487..fde6fffe92f4 100644
--- a/.github/workflows/ci-testing.yml
+++ b/.github/workflows/ci-testing.yml
@@ -133,3 +133,8 @@ jobs:
           python classify/predict.py --imgsz 32 --weights $b --source ../datasets/mnist2560/test/7/60.png  # predict
           python classify/predict.py --imgsz 32 --weights $m --source data/images/bus.jpg  # predict
           python export.py --weights $b --img 64 --imgsz 224 --include torchscript  # export
+          python - <<EOF
+          import torch
+          for path in '$m', '$b':
+              model = torch.hub.load('.', 'custom', path=path, source='local')
+          EOF
diff --git a/hubconf.py b/hubconf.py
index 011eaa57ff34..293f177dcbc1 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -30,7 +30,7 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
 
     from models.common import AutoShape, DetectMultiBackend
     from models.experimental import attempt_load
-    from models.yolo import Model
+    from models.yolo import DetectionModel
     from utils.downloads import attempt_download
     from utils.general import LOGGER, check_requirements, intersect_dicts, logging
     from utils.torch_utils import select_device
@@ -45,13 +45,13 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo
         if pretrained and channels == 3 and classes == 80:
             try:
                 model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
-                if autoshape:
+                if autoshape and isinstance(model.model, DetectionModel):
                     model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
             except Exception:
                 model = attempt_load(path, device=device, fuse=False)  # arbitrary model
         else:
             cfg = list((Path(__file__).parent / 'models').rglob(f'{path.stem}.yaml'))[0]  # model.yaml path
-            model = Model(cfg, channels, classes)  # create model
+            model = DetectionModel(cfg, channels, classes)  # create model
             if pretrained:
                 ckpt = torch.load(attempt_download(path), map_location=device)  # load
                 csd = ckpt['model'].float().state_dict()  # checkpoint state_dict as FP32