Export onnx protobuf bindings to python (#6651)

bddppq · web-flow · commit c43c91166260 · 2018-04-17T16:38:57.000-07:00
* Export onnx protobuf bindings to python

* rename native onnx module to _onnx
diff --git a/setup.py b/setup.py
@@ -686,6 +686,7 @@ def run(self):
     "torch/csrc/tensor/python_tensor.cpp",
     "torch/csrc/onnx/onnx.pb.cpp",
     "torch/csrc/onnx/onnx.cpp",
+    "torch/csrc/onnx/init.cpp",
 ]
 
 try:
diff --git a/torch/csrc/Module.cpp b/torch/csrc/Module.cpp
@@ -34,6 +34,7 @@
 #include "torch/csrc/jit/python_tracer.h"
 #include "torch/csrc/jit/init.h"
 #include "torch/csrc/jit/python_ir.h"
+#include "torch/csrc/onnx/init.h"
 
 #ifdef WITH_CUDNN
 #include "cudnn.h"
@@ -478,6 +479,7 @@ static PyObject* initModule() {
   ASSERT_TRUE(THPEngine_initModule(module));
   torch::autograd::initAutogradClosureBindings(module);
   torch::jit::initJITBindings(module);
+  torch::onnx::initONNXBindings(module);
   torch::autograd::initNNFunctions(module);
   torch::autograd::init_legacy_variable(module);
 #ifdef WITH_CUDA
diff --git a/torch/csrc/onnx/init.cpp b/torch/csrc/onnx/init.cpp
@@ -0,0 +1,26 @@
+#include "torch/csrc/onnx/init.h"
+#include "torch/csrc/onnx/onnx.pb.h"
+
+namespace torch { namespace onnx {
+void initONNXBindings(PyObject* module) {
+  auto m = py::handle(module).cast<py::module>();
+  auto onnx = m.def_submodule("_onnx");
+  py::enum_<onnx_TensorProto_DataType>(onnx, "TensorProtoDataType")
+      .value("UNDEFINED", onnx_TensorProto_DataType_UNDEFINED)
+      .value("FLOAT", onnx_TensorProto_DataType_FLOAT)
+      .value("UINT8", onnx_TensorProto_DataType_UINT8)
+      .value("INT8", onnx_TensorProto_DataType_INT8)
+      .value("UINT16", onnx_TensorProto_DataType_UINT16)
+      .value("INT16", onnx_TensorProto_DataType_INT16)
+      .value("INT32", onnx_TensorProto_DataType_INT32)
+      .value("INT64", onnx_TensorProto_DataType_INT64)
+      .value("STRING", onnx_TensorProto_DataType_STRING)
+      .value("BOOL", onnx_TensorProto_DataType_BOOL)
+      .value("FLOAT16", onnx_TensorProto_DataType_FLOAT16)
+      .value("DOUBLE", onnx_TensorProto_DataType_DOUBLE)
+      .value("UINT32", onnx_TensorProto_DataType_UINT32)
+      .value("UINT64", onnx_TensorProto_DataType_UINT64)
+      .value("COMPLEX64", onnx_TensorProto_DataType_COMPLEX64)
+      .value("COMPLEX128", onnx_TensorProto_DataType_COMPLEX128);
+}
+}} // namespace torch::onnx
diff --git a/torch/csrc/onnx/init.h b/torch/csrc/onnx/init.h
@@ -0,0 +1,9 @@
+#pragma once
+
+#include "torch/csrc/utils/pybind.h"
+
+namespace torch { namespace onnx {
+
+void initONNXBindings(PyObject* module);
+
+}} // namespace torch::onnx
diff --git a/torch/onnx/__init__.py b/torch/onnx/__init__.py
@@ -1,6 +1,10 @@
 import functools
 import types
 
+import torch._C as _C
+
+TensorProtoDataType = _C._onnx.TensorProtoDataType
+
 ONNX_ARCHIVE_MODEL_PROTO_NAME = "__MODEL_PROTO"
 
 
diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py
@@ -631,14 +631,14 @@ def _unique(g, input, sorted, return_inverse):
 # TODO: remove these once we support Type's in the JIT IR and we can once again
 # use the unified toType operator
 cast_pytorch_to_onnx = {
-    'uint8_t': 2,
-    'int8_t': 3,
-    'double': 11,
-    'float': 1,
-    'Half': 10,
-    'int': 6,
-    'int64_t': 7,
-    'int16_t': 5,
+    'uint8_t': torch.onnx.TensorProtoDataType.UINT8,
+    'int8_t': torch.onnx.TensorProtoDataType.INT8,
+    'double': torch.onnx.TensorProtoDataType.DOUBLE,
+    'float': torch.onnx.TensorProtoDataType.FLOAT,
+    'Half': torch.onnx.TensorProtoDataType.FLOAT16,
+    'int': torch.onnx.TensorProtoDataType.INT32,
+    'int64_t': torch.onnx.TensorProtoDataType.INT64,
+    'int16_t': torch.onnx.TensorProtoDataType.INT16,
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -686,6 +686,7 @@ def run(self):`
`686`	`686`	`"torch/csrc/tensor/python_tensor.cpp",`
`687`	`687`	`"torch/csrc/onnx/onnx.pb.cpp",`
`688`	`688`	`"torch/csrc/onnx/onnx.cpp",`
	`689`	`+ "torch/csrc/onnx/init.cpp",`
`689`	`690`	`]`
`690`	`691`
`691`	`692`	`try:`