enable export with dynamic axes (onnx#402)

Signed-off-by: Ksenija Stanojevic <ksenija.stanojevic@gmail.com> Co-authored-by: Wenbing Li <10278425+wenbingl@users.noreply.github.com>
LEMON-LM · Dec 16, 2020 · 00d95ba · 00d95ba
1 parent e7cfefd
commit 00d95ba
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/vision/classification/shufflenet/ShufflenetV2-export.py b/vision/classification/shufflenet/ShufflenetV2-export.py
@@ -16,7 +16,7 @@
 ]
 data_dir = 'test_data_set_0'
 
-url, filename = ("https://github.com/pytorch/hub/raw/master/dog.jpg", "dog.jpg")
+url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
 urllib.request.urlretrieve(url, filename)
 
 
@@ -135,7 +135,8 @@ def shufflenetv2_test():
         model_dir, data_dir = save_model('shufflenetv2', model.cpu(), input_1, output_1,
                                          opset_version=10,
                                          input_names=['input'],
-                                         output_names=['output'])
+                                         output_names=['output'],
+                                         dynamic_axes={"input": {0: 'batch_size'}, "output": {0: 'batch_size'}})
 
         # Test exported model with TensorProto data saved in files
         inputs_flatten = flatten(input_1)
@@ -161,5 +162,10 @@ def shufflenetv2_test():
 
         inference(model_dir, inputs, outputs)
 
+        # Test model with different input
+        input_2 = torch.randn(6, 3, 224, 224)
+        output_2 = model(input_2)
+        inference(model_dir, input_2, output_2)
 
-shufflenetv2_test()
+
+shufflenetv2_test()