Update YOLO code to use onnx==1.9.0 (instead of 1.4.1)

jkjung-avt · May 5, 2022 · 43677bf · 43677bf
1 parent c6afc40
commit 43677bf
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 35 deletions.
diff --git a/README.md b/README.md
@@ -258,10 +258,10 @@ Assuming this repository has been cloned at "${HOME}/project/tensorrt_demos", fo
    $ ./install_pycuda.sh
    ```
 
-2. Install **version "1.4.1" (not the latest version)** of python3 **"onnx"** module.  Note that the "onnx" module would depend on "protobuf" as stated in the [Prerequisite](#prerequisite) section.  Reference: [information provided by NVIDIA](https://devtalk.nvidia.com/default/topic/1052153/jetson-nano/tensorrt-backend-for-onnx-on-jetson-nano/post/5347666/#5347666).
+2. Install **version "1.9.0"** of python3 **"onnx"** module.  Note that the "onnx" module would depend on "protobuf" as stated in the [Prerequisite](#prerequisite) section.
 
    ```shell
-   $ sudo pip3 install onnx==1.4.1
+   $ sudo pip3 install onnx==1.9.0
    ```
 
 3. Go to the "plugins/" subdirectory and build the "yolo_layer" plugin.  When done, a "libyolo_layer.so" would be generated.

diff --git a/yolo/requirements.txt b/yolo/requirements.txt
@@ -1,3 +1,3 @@
-numpy>=1.15.1
-onnx==1.4.1
-pycuda>=2017.1.1
+numpy
+onnx==1.9.0
+pycuda
diff --git a/yolo/yolo_to_onnx.py b/yolo/yolo_to_onnx.py
@@ -360,23 +360,28 @@ def generate_param_name(self, param_category, suffix):
         param_name = self.node_name + '_' + param_category + '_' + suffix
         return param_name
 
-class UpsampleParams(object):
-    #Helper class to store the scale parameter for an Upsample node.
+class ResizeParams(object):
+    #Helper class to store the scale parameter for an Resize node.
 
     def __init__(self, node_name, value):
-        """Constructor based on the base node name (e.g. 86_Upsample),
+        """Constructor based on the base node name (e.g. 86_Resize),
         and the value of the scale input tensor.
 
         Keyword arguments:
-        node_name -- base name of this YOLO Upsample layer
-        value -- the value of the scale input to the Upsample layer as a numpy array
+        node_name -- base name of this YOLO Resize layer
+        value -- the value of the scale input to the Resize layer as numpy array
         """
         self.node_name = node_name
         self.value = value
 
     def generate_param_name(self):
-        """Generates the scale parameter name for the Upsample node."""
-        param_name = self.node_name + '_' + 'scale'
+        """Generates the scale parameter name for the Resize node."""
+        param_name = self.node_name + '_' + "scale"
+        return param_name
+
+    def generate_roi_name(self):
+        """Generates the roi input name for the Resize node."""
+        param_name = self.node_name + '_' + "roi"
         return param_name
 
 class WeightLoader(object):
@@ -393,26 +398,35 @@ def __init__(self, weights_file_path):
         """
         self.weights_file = self._open_weights_file(weights_file_path)
 
-    def load_upsample_scales(self, upsample_params):
+    def load_resize_scales(self, resize_params):
         """Returns the initializers with the value of the scale input
-        tensor given by upsample_params.
+        tensor given by resize_params.
 
         Keyword argument:
-        upsample_params -- a UpsampleParams object
+        resize_params -- a ResizeParams object
         """
         initializer = list()
         inputs = list()
-        name = upsample_params.generate_param_name()
-        shape = upsample_params.value.shape
-        data = upsample_params.value
+        name = resize_params.generate_param_name()
+        shape = resize_params.value.shape
+        data = resize_params.value
         scale_init = helper.make_tensor(
             name, TensorProto.FLOAT, shape, data)
         scale_input = helper.make_tensor_value_info(
             name, TensorProto.FLOAT, shape)
         initializer.append(scale_init)
         inputs.append(scale_input)
-        return initializer, inputs
 
+        # In opset 11 an additional input named roi is required. Create a dummy tensor to satisfy this.
+        # It is a 1D tensor of size of the rank of the input (4)
+        rank = 4
+        roi_name = resize_params.generate_roi_name()
+        roi_input = helper.make_tensor_value_info(roi_name, TensorProto.FLOAT, [rank])
+        roi_init = helper.make_tensor(roi_name, TensorProto.FLOAT, [rank], [0,0,0,0])
+        initializer.append(roi_init)
+        inputs.append(roi_input)
+
+        return initializer, inputs
 
     def load_conv_weights(self, conv_params):
         """Returns the initializers with weights from the weights file and
@@ -572,7 +586,7 @@ def build_onnx_graph(
                 initializer.extend(initializer_layer)
                 inputs.extend(inputs_layer)
             elif layer_type == 'upsample':
-                initializer_layer, inputs_layer = weight_loader.load_upsample_scales(
+                initializer_layer, inputs_layer = weight_loader.load_resize_scales(
                     params)
                 initializer.extend(initializer_layer)
                 inputs.extend(inputs_layer)
@@ -614,7 +628,7 @@ def _make_onnx_node(self, layer_name, layer_dict):
             node_creators['maxpool'] = self._make_maxpool_node
             node_creators['shortcut'] = self._make_shortcut_node
             node_creators['route'] = self._make_route_node
-            node_creators['upsample'] = self._make_upsample_node
+            node_creators['upsample'] = self._make_resize_node
             node_creators['yolo'] = self._make_yolo_node
 
             if layer_type in node_creators.keys():
@@ -873,7 +887,7 @@ def _make_route_node(self, layer_name, layer_dict):
                 route_node = helper.make_node(
                     'Split',
                     axis=1,
-                    split=[channels] * groups,
+                    #split=[channels] * groups,  # not needed for opset 11
                     inputs=[route_node_specs.name],
                     outputs=outputs,
                     name=layer_name,
@@ -916,37 +930,42 @@ def _make_route_node(self, layer_name, layer_dict):
             self._nodes.append(route_node)
         return layer_name, channels
 
-    def _make_upsample_node(self, layer_name, layer_dict):
-        """Create an ONNX Upsample node with the properties from
+    def _make_resize_node(self, layer_name, layer_dict):
+        """Create an ONNX Resize node with the properties from
         the DarkNet-based graph.
 
         Keyword arguments:
         layer_name -- the layer's name (also the corresponding key in layer_configs)
         layer_dict -- a layer parameter dictionary (one element of layer_configs)
         """
-        upsample_factor = float(layer_dict['stride'])
-        # Create the scales array with node parameters
-        scales = np.array([1.0, 1.0, upsample_factor, upsample_factor]).astype(np.float32)
+        resize_scale_factors = float(layer_dict['stride'])
+        # Create the scale factor array with node parameters
+        scales=np.array([1.0, 1.0, resize_scale_factors, resize_scale_factors]).astype(np.float32)
         previous_node_specs = self._get_previous_node_specs()
         inputs = [previous_node_specs.name]
 
         channels = previous_node_specs.channels
         assert channels > 0
-        upsample_params = UpsampleParams(layer_name, scales)
-        scales_name = upsample_params.generate_param_name()
-        # For ONNX opset >= 9, the Upsample node takes the scales array
-        # as an input.
+        resize_params = ResizeParams(layer_name, scales)
+
+        # roi input is the second input, so append it before scales
+        roi_name = resize_params.generate_roi_name()
+        inputs.append(roi_name)
+
+        scales_name = resize_params.generate_param_name()
         inputs.append(scales_name)
 
-        upsample_node = helper.make_node(
-            'Upsample',
+        resize_node = helper.make_node(
+            'Resize',
+            coordinate_transformation_mode='asymmetric',
             mode='nearest',
+            nearest_mode='floor',
             inputs=inputs,
             outputs=[layer_name],
             name=layer_name,
         )
-        self._nodes.append(upsample_node)
-        self.param_dict[layer_name] = upsample_params
+        self._nodes.append(resize_node)
+        self.param_dict[layer_name] = resize_params
         return layer_name, channels
 
     def _make_maxpool_node(self, layer_name, layer_dict):