feat(hub): loading a custom model with torch.hub.load (Megvii-BaseDetection#1396)

feat(hub): loading a custom model with `torch.hub.load`

Author: nick-konovalchuk

hubconf.py

@@ -5,6 +5,8 @@
 Usage example:
     import torch
     model = torch.hub.load("Megvii-BaseDetection/YOLOX", "yolox_s")
+    model = torch.hub.load("Megvii-BaseDetection/YOLOX", "yolox_custom",
+                           exp_path="exp.py", ckpt_path="ckpt.pth")
 """
 dependencies = ["torch"]
 
@@ -16,4 +18,5 @@
     yolox_l,
     yolox_x,
     yolov3,
+    yolox_custom
 )

yolox/models/build.py

@@ -14,6 +14,7 @@
     "yolox_l",
     "yolox_x",
     "yolov3",
+    "yolox_custom"
 ]
 
 _CKPT_ROOT_URL = "https://github.com/Megvii-BaseDetection/YOLOX/releases/download"
@@ -28,16 +29,20 @@
 }
 
 
-def create_yolox_model(
-    name: str, pretrained: bool = True, num_classes: int = 80, device=None
-) -> nn.Module:
+def create_yolox_model(name: str, pretrained: bool = True, num_classes: int = 80, device=None,
+                       exp_path: str = None, ckpt_path: str = None) -> nn.Module:
     """creates and loads a YOLOX model
 
     Args:
-        name (str): name of model. for example, "yolox-s", "yolox-tiny".
+        name (str): name of model. for example, "yolox-s", "yolox-tiny" or "yolox_custom"
+        if you want to load your own model.
         pretrained (bool): load pretrained weights into the model. Default to True.
-        num_classes (int): number of model classes. Defalut to 80.
-        device (str): default device to for model. Defalut to None.
+        device (str): default device to for model. Default to None.
+        num_classes (int): number of model classes. Default to 80.
+        exp_path (str): path to your own experiment file. Required if name="yolox_custom"
+        ckpt_path (str): path to your own ckpt. Required if name="yolox_custom" and you want to
+            load a pretrained model
+
 
     Returns:
         YOLOX model (nn.Module)
@@ -48,44 +53,59 @@ def create_yolox_model(
         device = "cuda:0" if torch.cuda.is_available() else "cpu"
     device = torch.device(device)
 
-    assert name in _CKPT_FULL_PATH, f"user should use one of value in {_CKPT_FULL_PATH.keys()}"
-    exp: Exp = get_exp(exp_name=name)
-    exp.num_classes = num_classes
-    yolox_model = exp.get_model()
-    if pretrained and num_classes == 80:
-        weights_url = _CKPT_FULL_PATH[name]
-        ckpt = load_state_dict_from_url(weights_url, map_location="cpu")
-        if "model" in ckpt:
-            ckpt = ckpt["model"]
-        yolox_model.load_state_dict(ckpt)
+    assert name in _CKPT_FULL_PATH or name == "yolox_custom", \
+        f"user should use one of value in {_CKPT_FULL_PATH.keys()} or \"yolox_custom\""
+    if name in _CKPT_FULL_PATH:
+        exp: Exp = get_exp(exp_name=name)
+        exp.num_classes = num_classes
+        yolox_model = exp.get_model()
+        if pretrained and num_classes == 80:
+            weights_url = _CKPT_FULL_PATH[name]
+            ckpt = load_state_dict_from_url(weights_url, map_location="cpu")
+            if "model" in ckpt:
+                ckpt = ckpt["model"]
+            yolox_model.load_state_dict(ckpt)
+    else:
+        assert exp_path is not None, "for a \"yolox_custom\" model exp_path must be provided"
+        exp: Exp = get_exp(exp_file=exp_path)
+        yolox_model = exp.get_model()
+        if ckpt_path:
+            ckpt = torch.load(ckpt_path, map_location="cpu")
+            if "model" in ckpt:
+                ckpt = ckpt["model"]
+            yolox_model.load_state_dict(ckpt)
 
     yolox_model.to(device)
     return yolox_model
 
 
-def yolox_nano(pretrained=True, num_classes=80, device=None):
+def yolox_nano(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-nano", pretrained, num_classes, device)
 
 
-def yolox_tiny(pretrained=True, num_classes=80, device=None):
+def yolox_tiny(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-tiny", pretrained, num_classes, device)
 
 
-def yolox_s(pretrained=True, num_classes=80, device=None):
+def yolox_s(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-s", pretrained, num_classes, device)
 
 
-def yolox_m(pretrained=True, num_classes=80, device=None):
+def yolox_m(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-m", pretrained, num_classes, device)
 
 
-def yolox_l(pretrained=True, num_classes=80, device=None):
+def yolox_l(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-l", pretrained, num_classes, device)
 
 
-def yolox_x(pretrained=True, num_classes=80, device=None):
+def yolox_x(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
     return create_yolox_model("yolox-x", pretrained, num_classes, device)
 
 
-def yolov3(pretrained=True, num_classes=80, device=None):
-    return create_yolox_model("yolox-tiny", pretrained, num_classes, device)
+def yolov3(pretrained: bool = True, num_classes: int = 80, device: str = None) -> nn.Module:
+    return create_yolox_model("yolov3", pretrained, num_classes, device)
+
+
+def yolox_custom(ckpt_path: str = None, exp_path: str = None, device: str = None) -> nn.Module:
+    return create_yolox_model("yolox_custom", ckpt_path=ckpt_path, exp_path=exp_path, device=device)