feat: add onnx export for torchvision models

Author: Bycob

.gitignore

@@ -6,3 +6,4 @@ build/*
 model/
 models/
 tools/build-cpp-netlib/
+__pycache__

ci/devel-trt.Dockerfile

@@ -112,7 +112,7 @@ RUN for url in \
         ; do curl -L -s -o /tmp/p.deb $url && dpkg -i /tmp/p.deb && rm -rf /tmp/p.deb; done
 
 RUN python3 -m pip install --upgrade pip
-RUN python3 -m pip install torch
+RUN python3 -m pip install torch torchvision
 
 RUN apt clean -y
 ADD ci/gitconfig /etc/gitconfig

ci/devel.Dockerfile

@@ -115,7 +115,7 @@ RUN for url in \
         ; do curl -L -s -o /tmp/p.deb $url && dpkg -i /tmp/p.deb && rm -rf /tmp/p.deb; done
 
 RUN python3 -m pip install --upgrade pip
-RUN python3 -m pip install torch
+RUN python3 -m pip install torch torchvision
 
 RUN apt clean -y
 ADD ci/gitconfig /etc/gitconfig

tests/CMakeLists.txt

@@ -483,3 +483,8 @@ if (USE_SIMSEARCH)
   endif()
 
 endif()
+
+# Python tests
+add_test(NAME ut_python
+       WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}/tests"
+       COMMAND python3 -m unittest ut_python -v)

tests/ut_python/__init__.py

@@ -0,0 +1,2 @@
+import unittest
+from .ut_tools_torch import *

tests/ut_python/temp/.gitignore

@@ -0,0 +1 @@
+*

tests/ut_python/ut_tools_torch.py

@@ -0,0 +1,70 @@
+import unittest
+import os
+import subprocess
+import torch
+import torchvision
+
+_test_dir = os.path.dirname(__file__)
+_temp_dir = os.path.join(_test_dir, "temp")
+
+def get_detection_input(batch_size=1):
+    """
+    Sample input for detection models, usable for tracing or testing
+    """
+    return (
+            torch.rand(batch_size, 3, 224, 224),
+            torch.full((batch_size,), 0).long(),
+            torch.Tensor([1, 1, 200, 200]).repeat((batch_size, 1)),
+            torch.full((batch_size,), 1).long(),
+    )
+
+class TestTorchvisionExport(unittest.TestCase):
+
+    def setUp(self):
+        os.chdir(os.path.join(_test_dir, "../../tools/torch"))
+
+    def test_resnet50_export(self):
+        # Export model (not pretrained because we don't have permission for the cache)
+        subprocess.run(["python3", "trace_torchvision.py", "-vp", "resnet50", "-o", _temp_dir])
+        model_file = os.path.join(_temp_dir, "resnet50.pt")
+        self.assertTrue(os.path.exists(model_file), model_file)
+
+        # Export to onnx
+        subprocess.run(["python3", "trace_torchvision.py", "-vp", "resnet50", "-o", _temp_dir, "--to-onnx", "--weights", model_file])
+        onnx_file = os.path.join(_temp_dir, "resnet50.onnx")
+        self.assertTrue(os.path.exists(onnx_file), onnx_file)
+
+    def test_fasterrcnn_export(self):
+        # Export model (not pretrained because we don't have permission for the cache)
+        subprocess.run(["python3", "trace_torchvision.py", "-vp", "fasterrcnn_resnet50_fpn", "-o", _temp_dir])
+        model_file = os.path.join(_temp_dir, "fasterrcnn_resnet50_fpn-cls91.pt")
+        self.assertTrue(os.path.exists(model_file), model_file)
+
+        # Test inference
+        rfcnn = torch.jit.load(model_file)
+        rfcnn.train()
+        model_loss, model_preds = rfcnn(*get_detection_input())
+        self.assertTrue(model_loss > 0)
+
+        rfcnn.eval()
+        model_loss, model_preds = rfcnn(torch.rand(1, 3, 224, 224))
+        self.assertTrue("boxes" in model_preds[0])
+
+        # Export to onnx
+        subprocess.run(["python3", "trace_torchvision.py", "-vp", "fasterrcnn_resnet50_fpn", "-o", _temp_dir, "--to-onnx", "--weights", model_file])
+        onnx_file = os.path.join(_temp_dir, "fasterrcnn_resnet50_fpn-cls91.onnx")
+        self.assertTrue(os.path.exists(onnx_file), onnx_file)
+
+    def tearDown(self):
+        print("Removing all files in %s" % _temp_dir)
+        ignore=[".gitignore"]
+        for f in os.listdir(_temp_dir):
+            removed = os.path.join(_temp_dir, f)
+            if f in ignore:
+                print("Ignore %s" % removed)
+            else:
+                print("Remove %s" % removed)
+                os.remove(removed)
+
+if __name__ == '__main__':
+    unittest.main()

tools/torch/trace_torchvision.py

@@ -36,6 +36,8 @@
 parser.add_argument('--backbone', type=str, help="Backbone for detection models")
 parser.add_argument('--print-models', action='store_true', help="Print all the available models names and exit")
 parser.add_argument('--to-dd-native', action='store_true', help="Prepare the model so that the weights can be loaded on native model with dede")
+parser.add_argument('--to-onnx', action="store_true", help="If specified, export to onnx instead of jit.")
+parser.add_argument('--weights', type=str, help="If not None, these weights will be embedded in the model before exporting")
 parser.add_argument('-a', "--all", action='store_true', help="Export all available models")
 parser.add_argument('-v', "--verbose", action='store_true', help="Set logging level to INFO")
 parser.add_argument('-o', "--output-dir", default=".", type=str, help="Output directory for traced models")
@@ -44,6 +46,9 @@
 parser.add_argument('--cpu', action='store_true', help="Force models to be exported for CPU device")
 parser.add_argument('--num_classes', type=int, help="Number of classes")
 parser.add_argument('--trace', action='store_true', help="Whether to trace model instead of scripting")
+parser.add_argument('--batch_size', type=int, default=1, help="When exporting with fixed batch size, this will be the batch size of the model")
+parser.add_argument('--img_width', type=int, default=224, help="Width of the image when exporting with fixed image size")
+parser.add_argument('--img_height', type=int, default=224, help="Height of the image when exporting with fixed image size")
 
 args = parser.parse_args()
 
@@ -112,15 +117,43 @@ def forward(self, x, ids = None, bboxes = None, labels = None):
 
         return loss, predictions
 
-def get_detection_input():
+
+class DetectionModel_PredictOnly(torch.nn.Module):
+    """
+    Adapt input and output of the model to make it exportable to
+    ONNX
+    """
+    def __init__(self, model):
+        super(DetectionModel_PredictOnly, self).__init__()
+        self.model = model
+
+    def forward(self, x):
+        l_x = [x[i] for i in range(x.shape[0])]
+        predictions = self.model(l_x)
+        # To dede format
+        pred_list = list()
+        for i in range(x.shape[0]):
+            pred_list.append(
+                    torch.cat((
+                        torch.full(predictions[i]["labels"].shape, i, dtype=float).unsqueeze(1),
+                        predictions[i]["labels"].unsqueeze(1).float(),
+                        predictions[i]["scores"].unsqueeze(1),
+                        predictions[i]["boxes"]), dim=1))
+
+        return torch.cat(pred_list)
+
+def get_image_input(batch_size=1, img_width=224, img_height=224):
+    return torch.rand(batch_size, 3, img_width, img_height)
+
+def get_detection_input(batch_size=1, img_width=224, img_height=224):
     """
     Sample input for detection models, usable for tracing or testing
     """
     return (
-            torch.rand(1, 3, 224, 224),
-            torch.full((1,), 0).long(),
-            torch.Tensor([1, 1, 200, 200]).unsqueeze(0),
-            torch.full((1,), 1).long(),
+            torch.rand(batch_size, 3, img_width, img_height),
+            torch.arange(0, batch_size).long(),
+            torch.Tensor([1, 1, 200, 200]).repeat((batch_size, 1)),
+            torch.full((batch_size,), 1).long(),
     )
 
 model_classes = {
@@ -230,7 +263,7 @@ def get_detection_input():
         else:
             if args.backbone:
                 raise RuntimeError("--backbone is only supported with models \"fasterrcnn\" or \"retinanet\".")
-            model = model_classes[mname](pretrained=args.pretrained, progress=args.verbose)
+            model = model_classes[mname](pretrained=args.pretrained, pretrained_backbone=args.pretrained, progress=args.verbose)
 
             if args.num_classes:
                 logging.info("Using num_classes = %d" % args.num_classes)
@@ -246,9 +279,17 @@ def get_detection_input():
                     # replace pretrained head
                     model.head = M.detection.retinanet.RetinaNetHead(in_channels, num_anchors, args.num_classes)
 
-        detect_model = DetectionModel(model)
-        detect_model.train()
-        script_module = torch.jit.script(detect_model)
+        if args.to_onnx:
+            model = DetectionModel_PredictOnly(model)
+            model.eval()
+        else:
+            model = DetectionModel(model)
+            model.train()
+            script_module = torch.jit.script(model)
+
+        if args.num_classes is None:
+            # TODO dont hard code this
+            args.num_classes = 91
 
     else:
         kwargs = {}
@@ -264,16 +305,45 @@ def get_detection_input():
 
         model.eval()
 
-
         # tracing or scripting model (default)
         if args.trace:
-            example = torch.rand(1, 3, 224, 224)
+            example = get_image_input(args.batch_size, args.img_width, args.img_height) 
             script_module = torch.jit.trace(model, example)
         else:
             script_module = torch.jit.script(model)
+
+    filename = os.path.join(
+            args.output_dir,
+            mname
+            + ("-pretrained" if args.pretrained else "")
+            + ("-" + args.backbone if args.backbone else "")
+            + ("-cls" + str(args.num_classes) if args.num_classes else "")
+            + ".pt")
+    
+    if args.weights:
+        # load weights
+        weights = torch.jit.load(args.weights).state_dict()
 
-    filename = os.path.join(args.output_dir, mname + ("-pretrained" if args.pretrained else "") + ("-" + args.backbone if args.backbone else "") + "-cls" + str(args.num_classes) + ".pt")
-    logging.info("Saving to %s", filename)
-    script_module.save(filename)
+        if args.to_onnx:
+            logging.info("Apply weights from %s to the onnx model" % args.weights)
+            model.load_state_dict(weights, strict=True)
+        else:
+            logging.info("Apply weights from %s to the jit model" % args.weights)
+            script_module.load_state_dict(weights, strict=True)
+
+    if args.to_onnx:
+        logging.info("Export model to onnx (%s)" % filename)
+        # remove extension
+        filename = filename[:-3] + ".onnx"
+        example = get_image_input(args.batch_size, args.img_width, args.img_height) 
+        torch.onnx.export(
+                model, example, filename,
+                export_params=True, verbose=args.verbose,
+                opset_version=11, do_constant_folding=True,
+                input_names=["input"], output_names=["output"])
+        # dynamic_axes={"input":{0:"batch_size"},"output":{0:"batch_size"}}
+    else:
+        logging.info("Saving to %s", filename)
+        script_module.save(filename)
 
 logging.info("Done")