Skip to content

Commit

Permalink
Add tests for networks with dynamic input size (parametrized)
Browse files Browse the repository at this point in the history
  • Loading branch information
sl-sergei committed Nov 17, 2020
1 parent d1242e6 commit b4d7f80
Show file tree
Hide file tree
Showing 28 changed files with 118 additions and 0 deletions.
Binary file not shown.
Binary file added testdata/dnn/onnx/data/input_gather_dynamic_axes.npy
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file added testdata/dnn/onnx/data/input_slice_dynamic_axes.npy
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file added testdata/dnn/onnx/data/output_slice_dynamic_axes.npy
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
97 changes: 97 additions & 0 deletions testdata/dnn/onnx/generate_onnx_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']])
Binary file not shown.
Binary file added testdata/dnn/onnx/models/gather_dynamic_axes.onnx
Binary file not shown.
Binary file not shown.
21 changes: 21 additions & 0 deletions testdata/dnn/onnx/models/maxpooling_sigmoid_dynamic_axes.onnx
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
pytorch1.6:�
V
01 MaxPool_0"MaxPool*
kernel_shape@@�*
pads@@@@�*
strides@@�

12 Sigmoid_1"Sigmoidtorch-jit-exportZ&
0!



height
widthb
2





B
Binary file not shown.
Binary file added testdata/dnn/onnx/models/slice_dynamic_axes.onnx
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.

0 comments on commit b4d7f80

Please sign in to comment.