Merge pull request #798 from sl-sergei:named_params_onnx

Author: alalek

testdata/dnn/onnx/data/input_average_pooling_dynamic_axes.npy

@@ -1040,3 +1040,100 @@ def forward(self, x):
 x = Variable(torch.zeros([1, 2, 2]))
 model = GatherMultiOutput()
 save_data_and_model("gather_multi_output", x, model)
+
+def postprocess_model(model_path, inputs_shapes):
+    onnx_model = onnx.load(model_path)
+
+    def update_inputs_dims(model, input_dims):
+        """
+            This function updates the sizes of dimensions of the model's inputs to the values
+            provided in input_dims. if the dim value provided is negative, a unique dim_param
+            will be set for that dimension.
+        """
+        def update_dim(tensor, dim, i, j, dim_param_prefix):
+            dim_proto = tensor.type.tensor_type.shape.dim[j]
+            if isinstance(dim, int):
+                if dim >= 0:
+                    dim_proto.dim_value = dim
+                else:
+                    dim_proto.dim_param = dim_param_prefix + str(i) + '_' + str(j)
+            elif isinstance(dim, str):
+                dim_proto.dim_param = dim
+            else:
+                raise ValueError('Only int or str is accepted as dimension value, incorrect type: {}'.format(type(dim)))
+
+        for i, input_dim_arr in enumerate(input_dims):
+            for j, dim in enumerate(input_dim_arr):
+                update_dim(model.graph.input[i], dim, i, j, 'in_')
+
+        onnx.checker.check_model(model)
+        return model
+    
+    onnx_model = update_inputs_dims(onnx_model, inputs_shapes)
+    onnx.save(onnx_model, model_path)
+
+class UnsqueezeAndConv(nn.Module):
+    def __init__(self):
+        super(UnsqueezeAndConv, self).__init__()
+        self.conv = nn.Conv2d(3, 3, kernel_size=1, stride=1, padding=0)
+    def forward(self, x):
+        x = x.unsqueeze(axis=0)
+        out = self.conv(x)
+        return out
+
+x = Variable(torch.randn(3, 10, 10))
+model = UnsqueezeAndConv()
+save_data_and_model("unsqueeze_and_conv_dynamic_axes", x, model)
+postprocess_model("models/unsqueeze_and_conv_dynamic_axes.onnx", [[3, 'height', 'width']])
+
+class SqueezeAndConv(nn.Module):
+    def __init__(self):
+        super(SqueezeAndConv, self).__init__()
+        self.conv = nn.Conv2d(3, 3, kernel_size=1, stride=1, padding=0)
+    def forward(self, x):
+        x = x.squeeze()
+        out = self.conv(x)
+        return out
+
+x = Variable(torch.randn(2, 1, 3, 3, 3))
+model = SqueezeAndConv()
+save_data_and_model("squeeze_and_conv_dynamic_axes", x, model)
+postprocess_model("models/squeeze_and_conv_dynamic_axes.onnx", [["batch_size", 1, "channels", 'height', 'width']])
+
+x = Variable(torch.randn(2))
+model = GatherScalar()
+save_data_and_model("gather_scalar_dynamic_axes", x, model)
+postprocess_model("models/gather_scalar_dynamic_axes.onnx", [['shape']])
+
+x = Variable(torch.randn(2, 2, 2, 2))
+print(x)
+model = Gather()
+print(model(x))
+print(model(x).shape)
+save_data_and_model("gather_dynamic_axes", x, model)
+postprocess_model("models/gather_dynamic_axes.onnx", [["batch_size", 2, 'height', 'width']])
+
+input = Variable(torch.randn(1, 2, 4, 4))
+model = Slice()
+save_data_and_model("slice_dynamic_axes", input, model)
+save_data_and_model("slice_opset_11_dynamic_axes", input, model, version=11)
+postprocess_model("models/slice_dynamic_axes.onnx", [["batch_size", 2, 'height', 'width']])
+postprocess_model("models/slice_opset_11_dynamic_axes.onnx", [["batch_size", 2, 'height', 'width']])
+
+x = Variable(torch.rand(1, 2, 2, 2))
+model = ResizeConv(2, 0, 2)
+save_data_and_model("resize_opset11_torch1.6_dynamic_axes", x, model, 11)
+postprocess_model("models/resize_opset11_torch1.6_dynamic_axes.onnx", [["batch_size", 2, 'height', 'width']])
+
+maxpooling_sigmoid = nn.Sequential(
+          nn.MaxPool2d(kernel_size=4, stride=2, padding=(1, 2), dilation=1),
+          nn.Sigmoid()
+        )
+input = Variable(torch.randn(2, 3, 12, 18))
+save_data_and_model("maxpooling_sigmoid_dynamic_axes", input, maxpooling_sigmoid)
+postprocess_model("models/maxpooling_sigmoid_dynamic_axes.onnx", [[2, 3, 'height', 'width']])
+
+ave_pool = nn.AvgPool2d(kernel_size=3, stride=2, padding=1)
+input = Variable(torch.randn(1, 3, 7, 5))
+save_data_and_model("average_pooling_dynamic_axes", input, ave_pool)
+postprocess_model("models/average_pooling_dynamic_axes.onnx", [[1, 3, 'height', 'width']])

testdata/dnn/onnx/data/input_gather_dynamic_axes.npy

@@ -0,0 +1,21 @@
+pytorch1.6:�

+V
+
0
1	MaxPool_0"MaxPool*
+kernel_shape@@�
*
+pads@
@@
@�
*
+strides@@�

+
+
1
2	Sigmoid_1"Sigmoidtorch-jit-exportZ&
+
0!
+

+
+
+height
+widthb
+
2
+

+
+
+
+
+B