Refactor ONNXRuntime Python interface

deployment/onnxruntime/README.md

@@ -28,7 +28,24 @@ The ONNXRuntime inference for `yolort`, both GPU and CPU are supported.
    cmake --build .
    ```
 
-1. Update your PyTorch model weights from ultralytics to yolort and export to ONNX following the [notebooks with tutorials](https://github.com/zhiqwang/yolov5-rt-stack/blob/master/notebooks/).
+1. Export your custom model to ONNX.
+
+   ```bash
+   python tools/export_model.py [--checkpoint_path path/to/custom/best.pt]
+                                [--simplify]
+   ```
+
+   Afterwards, you can see that a new pair of ONNX models ("best.onnx" and "best.sim.onnx") has been generated in the directory of "best.pt".
+
+1. \[Optional\] Quick test with the ONNXRuntime Python interface.
+
+   ```python
+   from yolort.runtime import PredictorORT
+
+   detector = PredictorORT("best.sim.onnx")
+   img_path = "bus.jpg"
+   scores, class_ids, boxes = detector.run_on_image(img_path)
+   ```
 
 1. Now, you can infer your own images.
 

tools/export_model.py

@@ -10,55 +10,64 @@
 except ImportError:
     onnxsim = None
 
-from yolort import models
+from yolort.models import YOLOv5
 
 
 def get_parser():
-    parser = argparse.ArgumentParser(
-        "CLI tool for exporting ONNX models", add_help=True
-    )
+    parser = argparse.ArgumentParser("CLI tool for exporting models.", add_help=True)
 
     parser.add_argument(
         "--checkpoint_path",
         type=str,
         required=True,
         help="The path of checkpoint weights",
     )
-    # Model architecture
-    parser.add_argument(
-        "--arch",
-        choices=["yolov5s", "yolov5m", "yolov5l"],
-        default="yolov5s",
-        help="Model architecture to export",
-    )
-    parser.add_argument(
-        "--num_classes", default=80, type=int, help="The number of classes"
-    )
     parser.add_argument(
         "--score_thresh",
         default=0.25,
         type=float,
         help="Score threshold used for postprocessing the detections.",
     )
     parser.add_argument(
-        "--export_friendly", action="store_true", help="Replace torch.nn.silu with Silu"
+        "--export_friendly",
+        action="store_true",
+        help="Replace torch.nn.SiLU with SiLU.",
     )
     parser.add_argument(
         "--image_size",
         default=640,
         type=int,
-        help="Image size for evaluation (default: 640)",
+        help="Image size for evaluation (default: 640).",
     )
-    parser.add_argument("--batch_size", default=1, type=int, help="Batch size")
+    parser.add_argument("--batch_size", default=1, type=int, help="Batch size.")
     parser.add_argument(
         "--opset", default=DEFAULT_OPSET, type=int, help="opset_version"
     )
-    parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model")
+    parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model.")
 
     return parser
 
 
-def export_onnx(model, inputs, export_onnx_path, opset_version, enable_simplify):
+def export_onnx(
+    model,
+    inputs,
+    export_onnx_path,
+    opset_version=11,
+    enable_simplify=False,
+):
+    """
+    Export the yolort models.
+
+    Args:
+        model (nn.Module): The model to be exported.
+        inputs (Tuple[torch.Tensor]): The inputs to the model.
+        export_onnx_path (str): A string containg a file name. A binary Protobuf
+            will be written to this file.
+        opset_version (int, default is 11): By default we export the model to the
+            opset version of the onnx submodule.
+        enable_simplify (bool, default is False): Whether to enable simplification
+            of the ONNX model.
+    """
     torch.onnx.export(
         model,
         inputs,
@@ -76,51 +85,55 @@ def export_onnx(model, inputs, export_onnx_path, opset_version, enable_simplify)
     )
 
     if enable_simplify:
-        export_onnx_sim_path = export_onnx_path.with_suffix(".sim.onnx")
-        if onnxsim is None:
-            raise ImportError("onnx-simplifier not found and is required by yolort")
+        input_shapes = {"images_tensors": list(inputs[0][0].shape)}
+        simplify_onnx(export_onnx_path, input_shapes)
+
 
-        print(f"Simplifing with onnx-simplifier {onnxsim.__version__}...")
+def simplify_onnx(onnx_path, input_shapes):
+    if onnxsim is None:
+        raise ImportError("onnx-simplifier not found and is required by yolort.")
 
-        # load onnx mode
-        onnx_model = onnx.load(export_onnx_path)
+    print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
 
-        # conver mode
-        model_sim, check = onnxsim.simplify(
-            onnx_model,
-            input_shapes={"images_tensors": list(inputs[0][0].shape)},
-            dynamic_input_shape=True,
-        )
+    # Load onnx mode
+    onnx_model = onnx.load(onnx_path)
 
-        assert check, "Simplified ONNX model could not be validated"
+    # Simlify the ONNX model
+    model_sim, check = onnxsim.simplify(
+        onnx_model,
+        input_shapes=input_shapes,
+        dynamic_input_shape=True,
+    )
 
-        onnx.save(model_sim, export_onnx_sim_path)
+    assert check, "There is something error when simplifying ONNX model"
+    export_onnx_sim_path = onnx_path.with_suffix(".sim.onnx")
+    onnx.save(model_sim, export_onnx_sim_path)
 
 
 def cli_main():
     parser = get_parser()
     args = parser.parse_args()
-    print("Command Line Args: {}".format(args))
+    print(f"Command Line Args: {args}")
     checkpoint_path = Path(args.checkpoint_path)
-    assert checkpoint_path.is_file(), f"Not found checkpoint: {checkpoint_path}"
+    assert checkpoint_path.exists(), f"Not found checkpoint file at '{checkpoint_path}'"
 
     # input data
-    images = torch.rand(3, args.image_size, args.image_size)
-    inputs = ([images],)
+    images = [torch.rand(3, args.image_size, args.image_size)]
+    inputs = (images,)
 
-    model = models.__dict__[args.arch](
-        num_classes=args.num_classes,
-        export_friendly=args.export_friendly,
-        score_thresh=args.score_thresh,
-    )
-    model.load_from_yolov5(checkpoint_path)
+    model = YOLOv5.load_from_yolov5(checkpoint_path, score_thresh=args.score_thresh)
     model.eval()
 
     # export ONNX models
     export_onnx_path = checkpoint_path.with_suffix(".onnx")
-    opset_version = args.opset
-    enable_simplify = args.simplify
-    export_onnx(model, inputs, export_onnx_path, opset_version, enable_simplify)
+
+    export_onnx(
+        model,
+        inputs,
+        export_onnx_path,
+        opset_version=args.opset,
+        enable_simplify=args.simplify,
+    )
 
 
 if __name__ == "__main__":

yolort/runtime/__init__.py

@@ -0,0 +1,3 @@
+from .ort_modeling import PredictorORT
+
+__all__ = ["PredictorORT"]