Adding support for exporting models with variable length input/output to ONNX (#20034)

pk-g · facebook-github-bot · commit 98e3aaeb78bc · 2019-06-05T12:02:23.000-07:00
Summary: Proposal: https://gist.github.com/pk-g/cc45ff8c5891b5699bffd883a87f13ae?fbclid=IwAR17bRA7Fks4APoZRYiNa93UkLdoFCpRDuIYEx0lNVyPTyaDAShbEnytiQo Pull Request resolved: #20034 Reviewed By: zrphercule Differential Revision: D15606731 Pulled By: houseroad fbshipit-source-id: 247251e07b4893cb3f7a1287948b1f57aadb7851
diff --git a/test/onnx/test_operators.py b/test/onnx/test_operators.py
@@ -248,6 +248,29 @@ def test_conv(self):
         x = torch.ones(20, 16, 50, 40, requires_grad=True)
         self.assertONNX(nn.Conv2d(16, 13, 3, bias=False), x)
 
+    def test_conv_variable_length(self):
+        x = torch.ones(5, 3, 6, 6, requires_grad=True)
+        model = torch.nn.Conv2d(3, 2, 3)
+        y = model(x)
+
+        dynamic_axes = {'input_1': [0, 2, 3], 'output_1': {0: 'output_1_variable_dim_0', 1: 'output_1_variable_dim_1'}}
+        model_proto_name = 'conv2d.onnx'
+        torch.onnx.export(model, x, model_proto_name, verbose=True, input_names=["input_1"], output_names=["output_1"],
+                          example_outputs=y, dynamic_axes=dynamic_axes)
+
+        import onnx
+        onnx_model = onnx.load(model_proto_name)
+        onnx.checker.check_model(onnx_model)
+
+        # Asserting the default dynamic axes names are generated when custom names are not provided
+        assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param == "input_1_dynamic_axes_1")
+        assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[2].dim_param == "input_1_dynamic_axes_2")
+        assert(onnx_model.graph.input[0].type.tensor_type.shape.dim[3].dim_param == "input_1_dynamic_axes_3")
+
+        # Asserting the custom names are applied when provided
+        assert(onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_param == "output_1_variable_dim_0")
+        assert(onnx_model.graph.output[0].type.tensor_type.shape.dim[1].dim_param == "output_1_variable_dim_1")
+
     def test_convtranspose(self):
         x = torch.ones(2, 3, 4, 5, requires_grad=True)
         self.assertONNX(nn.ConvTranspose2d(3, 3, 3, stride=3, bias=False,
diff --git a/torch/csrc/jit/export.cpp b/torch/csrc/jit/export.cpp
@@ -129,13 +129,17 @@ class EncoderBase {
       onnx::GraphProto* graph_proto,
       const std::shared_ptr<Graph>& graph,
       const std::map<std::string, at::Tensor>& initializers =
-          std::map<std::string, at::Tensor>());
+        std::map<std::string, at::Tensor>(),
+      const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes = 
+        std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>());
 
   void EncodeBlock(
       onnx::GraphProto* graph_proto,
       const Block* block,
       const std::map<std::string, at::Tensor>& initializers =
-          std::map<std::string, at::Tensor>());
+        std::map<std::string, at::Tensor>(),
+      const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes = 
+        std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>());
 
   virtual void EncodeTensor(
       onnx::TensorProto* tensor_proto,
@@ -149,7 +153,9 @@ class EncoderBase {
   virtual void EncodeValueInfo(
       onnx::GraphProto* graph_proto,
       onnx::ValueInfoProto* v,
-      const Value* n);
+      const Value* n,
+      const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes = 
+        std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>());
 
   void AddAttribute(
       onnx::NodeProto* node_proto,
@@ -206,16 +212,24 @@ EncoderBase::EncoderBase(
 void EncoderBase::EncodeValueInfo(
     onnx::GraphProto* graph_proto,
     onnx::ValueInfoProto* v,
-    const Value* n) {
-  v->set_name(n->uniqueName());
+    const Value* n,
+    const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes) {
+  std::string name = n->uniqueName();
+  v->set_name(name);
   if (CompleteTensorTypePtr node_type = n->type()->cast<CompleteTensorType>()) {
     onnx::TypeProto* t = v->mutable_type();
     onnx::TypeProto_Tensor* tensor_type = t->mutable_tensor_type();
     onnx::TensorShapeProto* shape = tensor_type->mutable_shape();
     const std::vector<std::int64_t>& sizes = node_type->sizes();
     for (size_t i = 0; i < sizes.size(); i++) {
       shape->add_dim();
-      shape->mutable_dim(i)->set_dim_value(sizes[i]);
+      if ((dynamic_axes.find(name) != dynamic_axes.end()) &&
+          (dynamic_axes.at(name).find(i) != dynamic_axes.at(name).end())){
+        shape->mutable_dim(i)->set_dim_param(dynamic_axes.at(name).at(i));
+      }
+      else{
+        shape->mutable_dim(i)->set_dim_value(sizes[i]);
+      }
     }
     tensor_type->set_elem_type(ATenTypeToOnnxType(node_type->scalarType()));
   } else if (BoolTypePtr node_type = n->type()->cast<BoolType>()) {
@@ -228,14 +242,16 @@ void EncoderBase::EncodeValueInfo(
 void EncoderBase::EncodeGraph(
     onnx::GraphProto* graph_proto,
     const std::shared_ptr<Graph>& graph,
-    const std::map<std::string, at::Tensor>& initializers) {
-  EncodeBlock(graph_proto, graph->block(), initializers);
+    const std::map<std::string, at::Tensor>& initializers,
+    const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes) {
+  EncodeBlock(graph_proto, graph->block(), initializers, dynamic_axes);
 }
 
 void EncoderBase::EncodeBlock(
     onnx::GraphProto* graph_proto,
     const Block* block,
-    const std::map<std::string, at::Tensor>& initializers) {
+    const std::map<std::string, at::Tensor>& initializers,
+    const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes) {
   AT_ASSERT(graph_proto != nullptr);
   std::string block_name = "torch-jit-export";
   if (num_blocks_) {
@@ -246,11 +262,11 @@ void EncoderBase::EncodeBlock(
 
   for (auto input : block->inputs()) {
     onnx::ValueInfoProto* v = graph_proto->add_input();
-    EncodeValueInfo(graph_proto, v, input);
+    EncodeValueInfo(graph_proto, v, input, dynamic_axes);
   }
   for (auto output : block->outputs()) {
     onnx::ValueInfoProto* v = graph_proto->add_output();
-    EncodeValueInfo(graph_proto, v, output);
+    EncodeValueInfo(graph_proto, v, output, dynamic_axes);
   }
   for (auto node : block->nodes()) {
     bool is_raw_export =
@@ -404,6 +420,7 @@ class GraphEncoder : public EncoderBase {
       int64_t onnx_opset_version,
       onnx_torch::OperatorExportTypes operator_export_type,
       const std::map<std::string, at::Tensor>& initializers,
+      const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes,
       bool defer_weight_export,
       bool strip_doc);
 
@@ -426,6 +443,7 @@ GraphEncoder::GraphEncoder(
     int64_t onnx_opset_version,
     onnx_torch::OperatorExportTypes operator_export_type,
     const std::map<std::string, at::Tensor>& initializers,
+    const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes,
     bool defer_weight_export,
     bool strip_doc)
     : EncoderBase(operator_export_type, strip_doc),
@@ -438,7 +456,7 @@ GraphEncoder::GraphEncoder(
   // This is the version of ONNX operator set we are targeting
   imp->set_version(onnx_opset_version);
 
-  EncodeGraph(model_proto_.mutable_graph(), graph, initializers);
+  EncodeGraph(model_proto_.mutable_graph(), graph, initializers, dynamic_axes);
 
   for (const std::string& domain : domains_) {
     auto* opset = model_proto_.add_opset_import();
@@ -1112,6 +1130,7 @@ std::string pretty_print_onnx(
       onnx_opset_version,
       operator_export_type,
       initializers,
+      std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>{},
       defer_weight_export,
       true);
   if (google_printer) {
@@ -1129,6 +1148,7 @@ std::tuple<std::string, RawDataExportMap> export_onnx(
     const std::shared_ptr<Graph>& graph,
     const std::map<std::string, at::Tensor>& initializers,
     int64_t onnx_opset_version,
+    const std::unordered_map<std::string, std::unordered_map<std::int64_t, std::string>>& dynamic_axes,
     bool defer_weight_export,
     ::torch::onnx::OperatorExportTypes operator_export_type,
     bool strip_doc_string) {
@@ -1137,6 +1157,7 @@ std::tuple<std::string, RawDataExportMap> export_onnx(
       onnx_opset_version,
       operator_export_type,
       initializers,
+      dynamic_axes,
       defer_weight_export,
       strip_doc_string);
   return std::make_tuple(
diff --git a/torch/csrc/jit/export.h b/torch/csrc/jit/export.h
@@ -25,6 +25,7 @@ TORCH_API std::tuple<std::string, RawDataExportMap> export_onnx(
     const std::shared_ptr<Graph>& graph,
     const std::map<std::string, at::Tensor>& initializers,
     int64_t onnx_opset_version,
+    const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes,
     bool defer_weight_export = false,
     ::torch::onnx::OperatorExportTypes operator_export_type =
         ::torch::onnx::OperatorExportTypes::ONNX,
diff --git a/torch/csrc/jit/python_ir.cpp b/torch/csrc/jit/python_ir.cpp
@@ -231,6 +231,7 @@ void initPythonIRBindings(PyObject* module_) {
           [](const std::shared_ptr<Graph> g,
              const std::map<std::string, at::Tensor>& initializers,
              int64_t onnx_opset_version,
+             const std::unordered_map<std::string, std::unordered_map<int64_t, std::string>>& dynamic_axes,
              bool defer_weight_export,
              ::torch::onnx::OperatorExportTypes operator_export_type,
              bool strip_doc_string) {
@@ -240,6 +241,7 @@ void initPythonIRBindings(PyObject* module_) {
                 g,
                 initializers,
                 onnx_opset_version,
+                dynamic_axes,
                 defer_weight_export,
                 operator_export_type,
                 strip_doc_string);
@@ -259,6 +261,7 @@ void initPythonIRBindings(PyObject* module_) {
           },
           py::arg("initializers"),
           py::arg("onnx_opset_version") = 0,
+          py::arg("dynamic_axes"),
           py::arg("defer_weight_export") = false,
           py::arg("operator_export_type") =
               ::torch::onnx::OperatorExportTypes::ONNX,
diff --git a/torch/onnx/utils.py b/torch/onnx/utils.py
@@ -56,7 +56,7 @@ def set_training(model, mode):
 def export(model, args, f, export_params=True, verbose=False, training=False,
            input_names=None, output_names=None, aten=False, export_raw_ir=False,
            operator_export_type=None, opset_version=None, _retain_param_name=True,
-           do_constant_folding=False, example_outputs=None, strip_doc_string=True):
+           do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None):
     r"""
     Export a model into ONNX format.  This exporter runs your model
     once in order to get a trace of its execution to be exported;
@@ -117,6 +117,41 @@ def export(model, args, f, export_params=True, verbose=False, training=False,
         strip_doc_string (bool, default True): if True, strips the field
             "doc_string" from the exported model, which information about the stack
             trace.
+        example_outputs: example outputs of the model that is being exported.
+        dynamic_axes (dict<string, dict<int, string>> or dict<string, list(int)>, default empty dict): 
+            a dictionary to specify dynamic axes of input/output, such that:            
+            - KEY:  input and/or output names
+            - VALUE: index of dynamic axes for given key and potentially the name to be used for
+            exported dynamic axes. In general the value is defined according to one of the following
+            ways or a combination of both:
+
+            (1). A list of integers specifiying the dynamic axes of provided input. In this scenario
+            automated names will be generated and applied to dynamic axes of provided input/output 
+            during export. 
+
+            OR (2). An inner dictionary that specifies a mapping FROM the index of dynamic axis in 
+            corresponding input/output TO the name that is desired to be applied on such axis of 
+            such input/output during export.
+
+            Example. if we have the following shape for inputs and outputs:             
+                shape(input_1) = ('b', 3, 'w', 'h')
+                and shape(input_2) = ('b', 4)           
+                and shape(output)  = ('b', 'd', 5)      
+
+            Then dynamic axes can be defined either as:
+            (a). ONLY INDICES:
+                 dynamic_axes = {'input_1':[0, 2, 3], 'input_2':[0], 'output':[0, 1]}
+                 where automatic names will be generated for exported dynamic axes
+
+            OR (b). INDICES WITH CORRESPONDING NAMES: 
+                dynamic_axes = {'input_1':{0:'batch', 1:'width', 2:'height'}, 
+                                 'input_2':{0:'batch'},
+                                 'output':{0:'batch', 1:'detections'}
+                 where provided names will be applied to exported dynamic axes
+
+            OR (c). MIXED MODE OF (a) and (b)
+                 dynamic_axes = {'input_1':[0, 2, 3], 'input_2':{0:'batch'}, 'output':[0,1]}
+
     """
     if aten or export_raw_ir:
         assert operator_export_type is None
@@ -130,7 +165,7 @@ def export(model, args, f, export_params=True, verbose=False, training=False,
     _export(model, args, f, export_params, verbose, training, input_names, output_names,
             operator_export_type=operator_export_type, opset_version=opset_version,
             _retain_param_name=_retain_param_name, do_constant_folding=do_constant_folding,
-            example_outputs=example_outputs, strip_doc_string=strip_doc_string)
+            example_outputs=example_outputs, strip_doc_string=strip_doc_string, dynamic_axes=dynamic_axes)
 
 
 # ONNX can't handle constants that are lists of tensors, which can
@@ -349,7 +384,7 @@ def _export(model, args, f, export_params=True, verbose=False, training=False,
             input_names=None, output_names=None, operator_export_type=OperatorExportTypes.ONNX,
             export_type=ExportTypes.PROTOBUF_FILE, example_outputs=None, propagate=False,
             opset_version=None, _retain_param_name=False, do_constant_folding=False,
-            strip_doc_string=True):
+            strip_doc_string=True, dynamic_axes=None):
     global __IN_ONNX_EXPORT
     assert __IN_ONNX_EXPORT is False
     __IN_ONNX_EXPORT = True
@@ -366,11 +401,17 @@ def _export(model, args, f, export_params=True, verbose=False, training=False,
 
         # TODO: Don't allocate a in-memory string for the protobuf
         defer_weight_export = export_type is not ExportTypes.PROTOBUF_FILE
+        if dynamic_axes is None:
+            dynamic_axes = {}
+
+        _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)
+
         if export_params:
-            proto, export_map = graph._export_onnx(params_dict, opset_version, defer_weight_export, operator_export_type,
-                                                   strip_doc_string)
+            proto, export_map = graph._export_onnx(
+                params_dict, opset_version, dynamic_axes, defer_weight_export, operator_export_type, strip_doc_string)
         else:
-            proto, export_map = graph._export_onnx({}, opset_version, False, operator_export_type, strip_doc_string)
+            proto, export_map = graph._export_onnx(
+                {}, opset_version, dynamic_axes, False, operator_export_type, strip_doc_string)
 
         if export_type == ExportTypes.PROTOBUF_FILE:
             assert(len(export_map) == 0)
@@ -721,6 +762,45 @@ def register_custom_op_symbolic(symbolic_name, symbolic_fn, opset_version):
     import torch.onnx.symbolic_registry as sym_registry
     sym_registry.register_op(op_name, symbolic_fn, ns, opset_version)
 
+# This helper function ensures dynamic axes argument is following the expected format
+def _validate_dynamic_axes(dynamic_axes, model, input_names, output_names):
+    if len(dynamic_axes) == 0:
+        return
+
+    if(hasattr(model, 'graph')):
+        # Extracting set of valid input/output names that shall be used for dynamic_axes
+        if (input_names is None) or len(input_names) == 0:
+            input_names = [x.uniqueName() for x in model.graph.inputs()]
+        if (output_names is None) or len(output_names) == 0:
+            output_names = [y.uniqueName() for y in model.graph.outputs()]
+
+    valid_names = set()
+    if input_names is not None:
+        valid_names.add(x for x in input_names)
+    if output_names is not None:
+        valid_names.add(x for x in output_names)
+
+    # If dynamic axes are provided as a list rather than dictionary, they should
+    # first get converted to a dictionary in expected format. If desired axes names
+    # are not provided for dynamic axes, automatic names shall be generated for
+    # provided dynamic axes of specified input/output
+    for key, value in dynamic_axes.items():
+        if key not in valid_names:
+            warnings.warn("Provided key {} for dynamic axes is not a valid input/output name".format(key))
+        if isinstance(value, list):
+            warnings.warn('No names were found for specified dynamic axes of provided input.' 
+                          'Automatically generated names will be applied to each dynamic axes of input {}'.format(key))
+
+            value_dict = {}
+            for i, x in enumerate(value):
+                if not isinstance(x, int):
+                    raise ValueError("The type of axis index is expected to be an integer")
+                if x in value_dict:
+                    warnings.warn('Duplicate dynamic axis index {} was provided for input {}.'
+                                  .format(x, key))
+                else:
+                    value_dict[x] = str(key) + '_dynamic_axes_' + str(i + 1)
+            dynamic_axes[key] = value_dict
 
 torch._C.Graph.op = _graph_op
 torch._C.Graph.at = _graph_at
