Export modules in ir with google protobuf by li-roy · Pull Request #9746 · pytorch/pytorch

li-roy · 2018-07-24T01:48:24Z

No description provided.

dzhulgakov

good stuff!

test/test_jit.py

+        f.seek(0)
+        import zipfile
+        with zipfile.ZipFile(f, 'r', compression=zipfile.ZIP_STORED) as z:
+            self.assertExpected(str([file.filename for file in z.infolist()]))


torch/csrc/jit/export.cpp

+          "\n\nDefined at:\n" + getNodeStackTraceString(node))
+    IR_ELSE()
+      // Special error messages for certain types of operators
+      if (node->kind() == aten::expand) {


torch/csrc/jit/export.cpp

+      operator_export_type_(operator_export_type) {
+  model_proto->set_producer_name("pytorch");
+  model_proto->set_ir_version(3);
+  model_proto->set_producer_version("0.4");


torch/csrc/jit/export.cpp

+      defer_weight_export_(defer_weight_export),
+      operator_export_type_(operator_export_type) {
+  model_proto->set_producer_name("pytorch");
+  model_proto->set_ir_version(3);


torch/csrc/jit/export.cpp

+        auto g = attr->add_graphs();
+        EncodeGraph(g, v);
+      }
+      break;


torch/csrc/jit/export.cpp

+  return std::string("\n") + idt(indent);
+}
+
+void dump(const onnx::TensorProto& tensor, std::ostream& stream) {


torch/csrc/jit/import.cpp

-
-  return Reader<Graph_>::read(&istream);
+at::Tensor JitDecoder::buildTensor(const onnx_torch::TensorProto& tensor_proto) {
+  at::Tensor tensor = at::CPU(onnxTypeToATenType(tensor_proto.data_type())).tensor();


torch/onnx/utils.py

+        compression = zipfile.ZIP_DEFLATED \
+            if export_type == ExportTypes.COMPRESSED_ZIP_ARCHIVE \
+            else zipfile.ZIP_STORED
+        with zipfile.ZipFile(f, 'w', compression=compression) as z:


torch/onnx/utils.py

    return torch_out


+def _export_module(module, f, operator_export_type=OperatorExportTypes.ONNX, export_type=ExportTypes.PROTOBUF_FILE):


torch/onnx/utils.py

    return torch_out


+def _export_module(module, f, operator_export_type=OperatorExportTypes.ONNX, export_type=ExportTypes.PROTOBUF_FILE):


zdevito

This looks really good! I found a few places where there are bugs (around storage offsets and sharing), and it needs a lot more tests, the contents of which I described in a comment.

test/test_jit.py

+        import zipfile
+        with zipfile.ZipFile(f, 'r', compression=zipfile.ZIP_STORED) as z:
+            self.assertExpected(str([file.filename for file in z.infolist()]))
+


torch/csrc/jit/export.cpp

-void encodeBlock(onnx::GraphProto * p_g, Block *b,
-                const std::vector<at::Tensor> & initializers,
-                ExportContext *ctx, RawDataExportMap* raw_data_export_map);
+class JitEncoder {


torch/csrc/jit/export.cpp

-}
-
-void encodeTypeProtoTensorType(onnx::TypeProto_Tensor* tensor_type, Value* n) {
+void EncodeTypeProtoTensorType(onnx::TypeProto_Tensor* tensor_type, const Value* n) {


torch/csrc/jit/export.cpp

+  // Add a buffer to the raw_data_export_map for the caller to dump into an
+  // external data store. If external_ref is not specified, we instead dump
+  // the contiguous data into the protobuf itself
+  if (defer_weight_export_) {


torch/csrc/jit/export.cpp

+                    const std::unique_ptr<script::Method> &method,
+                    const std::string prefix);
+
+  void EncodeTensor(onnx::TensorProto *tensor_proto,


torch/csrc/jit/import.cpp

+
+  std::shared_ptr<script::Module> curr = root_module;
+  for (size_t i = 0; i < vec.size() - 1; i++) {
+    if (curr->find_module(vec[i]) != nullptr) {


torch/csrc/jit/import.cpp

+    std::tie(parent_module, name) = parseFullName(root_module, tensor_proto.name());
+
+    auto param = buildParameter(tensor_proto);
+    parent_module->register_parameter(name, param, tensor_proto.int64_data(1));


torch/csrc/jit/import.h


 TORCH_API std::shared_ptr<Graph> ImportIRGraph(const std::string& serialized_graph, std::vector<at::Tensor> & initializers);

+std::shared_ptr<script::Module> ImportIRModule(


torch/onnx/utils.py

    return torch_out


+def _export_module(module, f, operator_export_type=OperatorExportTypes.ONNX, export_type=ExportTypes.PROTOBUF_FILE):


torch/onnx/utils.py

+        compression = zipfile.ZIP_DEFLATED \
+            if export_type == ExportTypes.COMPRESSED_ZIP_ARCHIVE \
+            else zipfile.ZIP_STORED
+        with zipfile.ZipFile(f, 'w', compression=compression) as z:


ezyang · 2018-07-27T04:32:27Z

@pytorchbot retest this please

ailzhang · 2018-07-31T16:45:24Z

@pytorchbot retest this please

zdevito

The changes to encode types look correct. I think this is good to go once we document the encoding of types, and the differences between parameter/tensor encodings.

torch/csrc/jit/export.cpp

torch/csrc/jit/import.cpp

facebook-github-bot