tests for Normalize subgraph, Slice, Mul and Expand

testdata/dnn/onnx/generate_onnx_models.py

@@ -8,6 +8,7 @@
 import numpy as np
 import os.path
 import onnx
+import onnxsim
 import google.protobuf.text_format
 import io
 
@@ -72,6 +73,14 @@ def save_onnx_data_and_model(input, output, name, operation, *args, **kwargs):
     model = onnx.helper.make_model(graph, producer_name=name)
     onnx.save(model, models_files)
 
+def simplify(name, rename=False, **kwargs):
+    model, check = onnxsim.simplify(name, **kwargs)
+    assert check, "couldn't valide"
+    name = name[:-5]
+    if rename:
+        name += '_optimized'
+    onnx.save(model, name + '.onnx')
+
 torch.manual_seed(0)
 np.random.seed(0)
 
@@ -412,6 +421,17 @@ def forward(self, x):
 save_data_and_model("slice", input, model)
 save_data_and_model("slice_opset_11", input, model, version=11)
 
+class SliceStarts(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SliceStarts, self).__init__()
+
+    def forward(self, x):
+        return x[-1:]
+
+model = SliceStarts()
+input_ = Variable(torch.randn(1, 10, dtype=torch.float32))
+save_data_and_model("slice_neg_starts", input_, model)
+
 input_2 = Variable(torch.randn(6, 6))
 custom_slice_list = [
     slice(1, 3, 1),
@@ -573,6 +593,18 @@ def forward(self, x):
 input_ = Variable(torch.tensor(list(range(20)), dtype=torch.float32))
 save_data_and_model("split_sizes", input_, model)
 
+class SplitAxis(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(SplitAxis, self).__init__()
+
+    def forward(self, x):
+        tup = torch.split(x, 2, -1)
+        return torch.cat(tup, 1)
+
+model = SplitAxis()
+input_ = Variable(torch.randn(1, 10, dtype=torch.float32))
+save_data_and_model("split_neg_axis", input_, model)
+
 class SplitMax(nn.Module):
 
     def __init__(self):
@@ -840,6 +872,32 @@ def forward(self, x):
 output = np.mean(x, axis=2, keepdims=True)
 save_onnx_data_and_model(x, output, 'reduce_mean_axis2', 'ReduceMean', axes=(2), keepdims=True)
 
+class Expand(nn.Module):
+    def __init__(self):
+        super(Expand, self).__init__()
+
+    def forward(self, x):
+        return x.expand(1, 3, -1, -1, -1)
+
+input = Variable(torch.randn(1, 3, 2, 4))
+model = Expand()
+model.eval()
+save_data_and_model("expand", input, model, export_params=True, version=12)
+simplify('models/expand.onnx', False)
+
+class ExpandIdentity(nn.Module):
+    def __init__(self):
+        super(ExpandIdentity, self).__init__()
+
+    def forward(self, x):
+        return x.expand(1, 3, -1, -1)
+
+input = Variable(torch.randn(1, 3, 2, 4))
+model = ExpandIdentity()
+model.eval()
+save_data_and_model("expand_identity", input, model, export_params=True, version=12)
+simplify('models/expand_identity.onnx', False)
+
 class Expand(nn.Module):
     def __init__(self, shape):
         super(Expand, self).__init__()
@@ -908,6 +966,23 @@ def forward(self, x):
 x = Variable(torch.randn(1, 2, 3, 4))
 save_data_and_model("reduceL2_subgraph_2", x, model)
 
+class reduceL2_subgraph2_2(nn.Module):
+    def __init__(self):
+        super(reduceL2_subgraph2_2, self).__init__()
+        self.size = torch.Size([1, 3, 2, 4])
+
+    def forward(self, x):
+        norm = torch.norm(x, p=2, dim=1, keepdim=True)
+        clip = torch.clamp(norm, min=0)
+        expand = clip.expand([1, 3, 2, 4])
+        return x / expand
+
+input = Variable(torch.randn(1, 3, 2, 4))
+model = reduceL2_subgraph2_2()
+model.eval()
+save_data_and_model("reduceL2_subgraph2_2", input, model, export_params=True, version=12)
+simplify('models/reduceL2_subgraph2_2.onnx', False)
+
 from torchvision.ops.misc import *
 n = 3
 model = FrozenBatchNorm2d(n)
@@ -1148,6 +1223,18 @@ def forward(self, x0, x1, x2):
 input_2 = Variable(torch.ones(2, 1, 4, 1, dtype=torch.float32))
 save_data_and_model_multy_inputs("scale_broadcast", model, input_0, input_1, input_2)
 
+class ScaleBroadcastMid(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super(ScaleBroadcastMid, self).__init__()
+
+    def forward(self, x0, x1):
+        return torch.mul(x0, x1)
+
+model = ScaleBroadcastMid()
+input_0 = Variable(torch.ones(2, 1, 4, dtype=torch.float32))
+input_1 = Variable(torch.ones(2, 5, 4, dtype=torch.float32))
+save_data_and_model_multy_inputs("scale_broadcast_mid", model, input_0, input_1)
+
 x = Variable(torch.randn(1, 3, 25))
 conv1d = nn.Conv1d(3, 2, kernel_size=3, padding=2, stride=2, dilation=2, bias=False)
 save_data_and_model("conv1d", x, conv1d)

testdata/dnn/onnx/models/scale_broadcast_mid.onnx

@@ -0,0 +1,19 @@
+pytorch1.9:t
+
+
0
+
1
2Mul_0"Multorch-jit-exportZ
+
0
+

+
+

+Z
+
1
+

+
+
+b
+
2
+

+
+
+B	

testdata/dnn/onnx/models/split_neg_axis.onnx

@@ -0,0 +1,22 @@
+pytorch1.9:�

+S
+tensor
1
2
3
4
5Split_0"Split*
+axis���������
�
*
+split@@@@@�

+1
+
1
+
2
+
3
+
4
+
5
6Concat_1"Concat*
+axis
�
torch-jit-exportZ
+tensor
+

+

+
+b
+
6
+

+

+
+B