[ONNX] Utilize ONNX shape inference for ONNX exporter (pytorch#40628)

Summary:
It is often that the conversion from torch operator to onnx operator requires input rank/dtype/shape to be known. Previously, the conversion depends on tracer to provide these info, leaving a gap in conversion of scripted modules.

We are extending the export with support from onnx shape inference. If enabled, onnx shape inference will be called whenever an onnx node is created. This is the first PR introducing the initial look of the feature. More and more cases will be supported following this PR.

* Added pass to run onnx shape inference on a given node. The node has to have namespace `onnx`.
* Moved helper functions from `export.cpp` to a common place for re-use.
* This feature is currently experimental, and can be turned on through flag `onnx_shape_inference` in internal api `torch.onnx._export`.
* Currently skipping ONNX Sequence ops, If/Loop and ConstantOfShape due to limitations. Support will be added in the future.

Pull Request resolved: pytorch#40628

Reviewed By: mrshenli

Differential Revision: D22709746

Pulled By: bzinodev

fbshipit-source-id: b52aeeae00667e66e0b0c1144022f7af9a8b2948

Author: BowenBao

.github/workflows/lint.yml

@@ -152,6 +152,9 @@ jobs:
             --verbose                                              \
             --paths torch/csrc/                                    \
             --diff "$MERGE_BASE"                                   \
+            -g"-torch/csrc/jit/passes/onnx/helper.cpp"             \
+            -g"-torch/csrc/jit/passes/onnx/shape_type_inference.cpp"\
+            -g"-torch/csrc/jit/serialization/onnx.cpp"             \
             -g"-torch/csrc/jit/serialization/export.cpp"           \
             -g"-torch/csrc/jit/serialization/import.cpp"           \
             -g"-torch/csrc/jit/serialization/import_legacy.cpp"    \

aten/src/ATen/core/interned_strings.h

@@ -283,6 +283,7 @@ namespace c10 {
   _(onnx, SequenceConstruct)         \
   _(onnx, SequenceEmpty)             \
   _(onnx, SequenceInsert)            \
+  _(onnx, SequenceErase)             \
   _(onnx, ConcatFromSequence)        \
   _(onnx, Identity)                  \
   _(onnx, SoftmaxCrossEntropyLoss)   \

caffe2/CMakeLists.txt

@@ -459,6 +459,7 @@ if(NOT INTERN_BUILD_MOBILE OR NOT BUILD_CAFFE2_MOBILE)
   if(NOT INTERN_BUILD_MOBILE)
     list(APPEND TORCH_SRCS
       ${TORCH_SRC_DIR}/csrc/api/src/jit.cpp
+      ${TORCH_SRC_DIR}/csrc/jit/serialization/onnx.cpp
       ${TORCH_SRC_DIR}/csrc/jit/serialization/export.cpp
       ${TORCH_SRC_DIR}/csrc/jit/serialization/export_module.cpp
       ${TORCH_SRC_DIR}/csrc/jit/serialization/import_legacy.cpp

test/onnx/test_pytorch_common.py

@@ -86,5 +86,14 @@ def wrapper(self):
         return wrapper
     return skip_dec
 
+def skipIfONNXShapeInference(onnx_shape_inference):
+    def skip_dec(func):
+        def wrapper(self):
+            if self.onnx_shape_inference is onnx_shape_inference:
+                raise unittest.SkipTest("Skip verify test for unsupported opset_version")
+            return func(self)
+        return wrapper
+    return skip_dec
+
 def flatten(x):
     return tuple(function._iter_filter(lambda o: isinstance(o, torch.Tensor))(x))

test/onnx/test_pytorch_onnx_onnxruntime.py

@@ -17,7 +17,7 @@
 from model_defs.rnn_model_with_packed_sequence import RnnModelWithPackedSequence
 from test_pytorch_common import (skipIfUnsupportedMinOpsetVersion, enableScriptTest,
                                  skipIfUnsupportedOpsetVersion, skipIfNoLapack,
-                                 skipIfUnsupportedMaxOpsetVersion)
+                                 skipIfUnsupportedMaxOpsetVersion, skipIfONNXShapeInference)
 from test_pytorch_common import BATCH_SIZE
 from test_pytorch_common import RNN_BATCH_SIZE, RNN_SEQUENCE_LENGTH, RNN_INPUT_SIZE, RNN_HIDDEN_SIZE
 import model_defs.word_language_model as word_language_model
@@ -79,7 +79,8 @@ def run_model_test(self, model, batch_size=2, state_dict=None,
                            keep_initializers_as_inputs=self.keep_initializers_as_inputs,
                            dynamic_axes=dynamic_axes,
                            input_names=input_names, output_names=output_names,
-                           fixed_batch_size=fixed_batch_size)
+                           fixed_batch_size=fixed_batch_size,
+                           onnx_shape_inference=self.onnx_shape_inference)
 
         # compute onnxruntime output prediction
         ort_sess = onnxruntime.InferenceSession(f.getvalue())
@@ -103,6 +104,7 @@ class TestONNXRuntime(unittest.TestCase):
     from torch.onnx.symbolic_helper import _export_onnx_opset_version
     opset_version = _export_onnx_opset_version
     keep_initializers_as_inputs = True  # For IR version 3 type export.
+    onnx_shape_inference = False
 
     def setUp(self):
         torch.manual_seed(0)
@@ -496,15 +498,15 @@ def test_tensor(self):
         class ScalarInputModel(torch.jit.ScriptModule):
             @torch.jit.script_method
             def forward(self, input):
-                return torch.tensor(input.shape[1]) 
+                return torch.tensor(input.shape[1])
 
         x = torch.randn(3, 4)
         self.run_test(ScalarInputModel(), x)
 
         class TensorInputModel(torch.jit.ScriptModule):
             @torch.jit.script_method
             def forward(self, input):
-                return torch.tensor([input.shape[0], input.shape[1]]) 
+                return torch.tensor([input.shape[0], input.shape[1]])
 
         x = torch.randn(3, 4)
         self.run_test(TensorInputModel(), x)
@@ -520,15 +522,15 @@ def forward(self, input):
         class InputWithDtypeModel(torch.jit.ScriptModule):
             @torch.jit.script_method
             def forward(self, input):
-                return torch.tensor(input.shape[1], dtype=torch.long) 
+                return torch.tensor(input.shape[1], dtype=torch.long)
 
         x = torch.randn(3, 4)
         self.run_test(InputWithDtypeModel(), x)
 
         class MixedInputModel(torch.jit.ScriptModule):
             @torch.jit.script_method
             def forward(self, input):
-                return torch.tensor([input.shape[0], int(input)]) 
+                return torch.tensor([input.shape[0], int(input)])
 
         x = torch.randn(1)
         self.run_test(MixedInputModel(), x)
@@ -686,6 +688,23 @@ def forward(self, input1, input2, input3):
         self.run_test(TraceModel(), (x1, x2, x3), atol=10e-5)
         self.run_test(ScriptModel(), (x1, x2, x3), atol=10e-5)
 
+    # Conversion of Transpose depends on input shape to be known.
+    # The following test only works when onnx shape inference is enabled.
+    @skipIfONNXShapeInference(False)
+    def test_transpose_infer_shape(self):
+        class TransposeModule(torch.jit.ScriptModule):
+            def __init__(self):
+                super(TransposeModule, self).__init__()
+                self.conv = torch.nn.Conv2d(3, 1, 3, stride=2)
+
+            @torch.jit.script_method
+            def forward(self, x):
+                x = self.conv(x)
+                return x.transpose(0, 1)
+
+        x = torch.randn(32, 3, 64, 64)
+        self.run_test(TransposeModule(), x)
+
     def squeeze_model_tests(self, d, x1, x2):
         class Squeeze(torch.nn.Module):
             def forward(self, x):
@@ -842,6 +861,23 @@ def forward(self, x):
         x = torch.randn(2, 3, 4)
         self.run_test(ArithmeticModule(), x)
 
+    # In scripting the first transpose node do not carry shape and dtype info.
+    # The following test only works when onnx shape inference is enabled.
+    @skipIfONNXShapeInference(False)
+    def test_arithmetic_infer_dtype(self):
+        class ArithmeticModule(torch.jit.ScriptModule):
+            @torch.jit.script_method
+            def forward(self, x):
+                x = x.t()
+                x = x + 2
+                x = x - 4
+                x = x * 6
+                x = x / 8
+                return x
+
+        x = torch.randn(2, 3)
+        self.run_test(ArithmeticModule(), x)
+
     def test_floor_div(self):
         class FloorDivModule(torch.nn.Module):
             def forward(self, x, y):
@@ -3015,6 +3051,21 @@ def forward(self, x):
         x = torch.randn(4, 2, 3, requires_grad=True)
         self.run_test(UnfoldModel(), x)
 
+    @skipIfONNXShapeInference(False)
+    def test_unfold_infer_shape(self):
+        class UnfoldModule(torch.jit.ScriptModule):
+            def __init__(self):
+                super(UnfoldModule, self).__init__()
+                self.conv = torch.nn.Conv1d(3, 1, 3, stride=2)
+
+            @torch.jit.script_method
+            def forward(self, x):
+                x = self.conv(x)
+                return x.unfold(dimension=2, size=2, step=2)
+
+        x = torch.randn(32, 3, 64)
+        self.run_test(UnfoldModule(), x)
+
     def test_remainder(self):
         class RemainderModel(torch.nn.Module):
             def forward(self, input, other):
@@ -4187,5 +4238,11 @@ def setup_rnn_tests():
                                          keep_initializers_as_inputs=False))
 
 
+# opset 12 tests, with _onnx_shape_inference=True.
+TestONNXRuntime_opset12_onnx_shape_inference = type(str("TestONNXRuntime_opset12_onnx_shape_inference"),
+                                                    (unittest.TestCase,),
+                                                    dict(TestONNXRuntime.__dict__, opset_version=12,
+                                                         onnx_shape_inference=True))
+
 if __name__ == '__main__':
     unittest.main()

tools/build_variables.bzl

@@ -322,6 +322,7 @@ libtorch_extra_sources = libtorch_core_jit_sources + [
     "torch/csrc/jit/mobile/observer.cpp",
     "torch/csrc/jit/mobile/optim/sgd.cpp",
     "torch/csrc/jit/mobile/sequential.cpp",
+    "torch/csrc/jit/serialization/onnx.cpp",
     "torch/csrc/jit/serialization/export.cpp",
     "torch/csrc/jit/serialization/export_module.cpp",
     "torch/csrc/jit/serialization/import_legacy.cpp",
@@ -501,6 +502,7 @@ libtorch_python_core_sources = [
     "torch/csrc/jit/passes/onnx/scalar_type_analysis.cpp",
     "torch/csrc/jit/passes/onnx/unpack_quantized_weights.cpp",
     "torch/csrc/jit/passes/onnx/prepare_inplace_ops_for_onnx.cpp",
+    "torch/csrc/jit/passes/onnx/shape_type_inference.cpp",
     "torch/csrc/jit/python/python_arg_flatten.cpp",
     "torch/csrc/jit/python/python_custom_class.cpp",
     "torch/csrc/jit/python/python_interpreter.cpp",

tools/git-pre-commit

@@ -10,6 +10,9 @@ then
   python tools/clang_tidy.py \
     --paths torch/csrc \
     --diff HEAD \
+    -g"-torch/csrc/jit/passes/onnx/helper.cpp" \
+    -g"-torch/csrc/jit/passes/onnx/shape_type_inference.cpp" \
+    -g"-torch/csrc/jit/serialization/onnx.cpp" \
     -g"-torch/csrc/jit/serialization/export.cpp" \
     -g"-torch/csrc/jit/serialization/import.cpp" \
     -j

torch/csrc/jit/passes/onnx/constant_fold.cpp

@@ -103,7 +103,10 @@ c10::optional<at::Tensor> runTorchSlice_opset9(
 c10::optional<at::Tensor> runTorchSlice_opset10(
     const Node* node,
     std::vector<at::Tensor>& inputTensorValues) {
-  if (inputTensorValues.size() < 3 || inputTensorValues.size() > 5) {
+  const int maxSliceInputCount = 5;
+  const int minSliceInputCount = 3;
+  if (inputTensorValues.size() < minSliceInputCount ||
+      inputTensorValues.size() > maxSliceInputCount) {
     std::cerr
         << "Warning: Constant folding - Invalid number of inputs found for opset 10 or 11 onnx::Slice op. "
         << "Constant folding not applied." << std::endl;
@@ -249,11 +252,9 @@ c10::optional<at::Tensor> runTorchBackendForOnnx(
     return c10::optional<at::Tensor>(updated_val);
   } else if (node->kind() == onnx::Cast) {
     assert(inputTensorValues.size() == 1);
-    if (node->hasAttributeS("to") &&
-        onnxTypeToScalarTypeMap.find(node->i(attr::to)) !=
-            onnxTypeToScalarTypeMap.end()) {
-      updated_val =
-          inputTensorValues[0].to(onnxTypeToScalarTypeMap[node->i(attr::to)]);
+    if (node->hasAttributeS("to") && ONNXTypeToATenType(node->i(attr::to))) {
+      updated_val = inputTensorValues[0].to(
+          ONNXTypeToATenType(node->i(attr::to)).value());
       return c10::optional<at::Tensor>(updated_val);
     }
     return c10::nullopt;

torch/csrc/jit/passes/onnx/helper.cpp

@@ -1,5 +1,5 @@
 #include <torch/csrc/jit/passes/onnx/helper.h>
-#include <torch/csrc/jit/jit_log.h>
+#include <onnx/onnx_pb.h>
 
 namespace torch {
 namespace jit {
@@ -59,5 +59,40 @@ Node* addNodeToBlock(Block* block, Value* input, Symbol kind) {
   }
   return new_node;
 }
+
+c10::optional<at::ScalarType> ONNXTypeToATenType(int32_t onnx_type) {
+  switch (onnx_type) {
+    case ::ONNX_NAMESPACE::TensorProto_DataType_UNDEFINED:
+      return at::ScalarType::Undefined;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT:
+      return at::kFloat;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_UINT8:
+      return at::kByte;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_INT8:
+      return at::kChar;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_INT16:
+      return at::kShort;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_INT32:
+      return at::kInt;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_INT64:
+      return at::kLong;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_BOOL:
+      return at::kBool;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_FLOAT16:
+      return at::kHalf;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_DOUBLE:
+      return at::kDouble;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_COMPLEX64:
+      return at::kComplexFloat;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_COMPLEX128:
+      return at::kComplexDouble;
+    case ::ONNX_NAMESPACE::TensorProto_DataType_BFLOAT16:
+      return at::kBFloat16;
+    default:
+      TORCH_CHECK("unexpected tensor scalar type");
+  }
+  return c10::optional<at::ScalarType>{};
+}
+
 } // namespace jit
 } // namespace torch

torch/csrc/jit/passes/onnx/helper.h

@@ -8,18 +8,19 @@
 namespace torch {
 namespace jit {
 
-namespace onnx {
 static const int OPSET_VERSION_1 = 1;
 static const int OPSET_VERSION_9 = 9;
 static const int OPSET_VERSION_10 = 10;
 static const int OPSET_VERSION_11 = 11;
 static const int OPSET_VERSION_12 = 12;
-} // namespace onnx
 
 using ValueToParamPairMap = std::map<Value*, std::pair<std::string, IValue>>;
 
 using ParamMap = std::map<std::string, IValue>;
 
+void buildParamsMapFromValueToParamsMap(
+    const ValueToParamPairMap& valsToParamsMap,
+    ParamMap& paramsDict);
 ValueToParamPairMap buildValueToParamsMap(Block* b, const ParamMap& paramsDict);
 void eraseUnusedValuesFromMap(ValueToParamPairMap& valsToParamsMap);
 void eraseUnusedBlockInputs(Block* b);
@@ -28,5 +29,6 @@ void buildParamsMapFromValueToParamsMap(
     ParamMap& paramsDict);
 Node* addNodeToBlock(Block* block, Value* input, Symbol kind);
 
+TORCH_API c10::optional<at::ScalarType> ONNXTypeToATenType(int32_t onnx_type);
 } // namespace jit
 } // namespace torch

torch/csrc/jit/passes/onnx/peephole.cpp

@@ -576,7 +576,7 @@ static void eraseListConstruct(Block* block, int opset_version) {
               i, std::vector<Value*>({concat_node->output()}));
 
         } else {
-          if (opset_version < onnx::OPSET_VERSION_11) {
+          if (opset_version < OPSET_VERSION_11) {
             // Tensor lists are used mostly for inputs to cat/stack. They are
             // already handled in those symbolics, and should become dead
             // afterwards.