Add save() to torch._C.Function (#20386)

davidriazati · facebook-github-bot · commit 410c7210db40 · 2019-05-20T16:19:51.000-07:00
Summary: Fixes #20017 This wraps the `torch._C.Function` currently returned from `torch.jit.script` and `torch.jit.trace` in a `ScriptFunction` and `TracedFunction` respectively, both of which are just wrappers to hold the function. ](https://our.intern.facebook.com/intern/diff/15403161/) Pull Request resolved: #20386 Pulled By: driazati Differential Revision: D15403161 fbshipit-source-id: 94fb9f32929e62a00be6cf7512ea144ec9b91e0b
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -3127,6 +3127,21 @@ def fn(x):
         warns = [str(w.message) for w in warns]
         self.assertEqual(len(warns), 0)
 
+    @unittest.skipIf(sys.platform == "win32", "temp file name on windows")
+    def test_trace_save(self):
+        def fn(x):
+            return x + 2
+
+        def check(func):
+            with tempfile.NamedTemporaryFile() as f:
+                func.save(f.name)
+                loaded = torch.jit.load(f.name)
+                input = torch.randn(2, 2)
+                self.assertEqual(func(input), loaded(input))
+
+        out = torch.jit.trace(fn, (torch.ones(2, 2),))
+        check(out)
+
     @unittest.skipIf(sys.platform == "win32", "TODO: need to fix this test case for Windows")
     def test_torch_load_error(self):
         class J(torch.jit.ScriptModule):
diff --git a/torch/csrc/jit/script/init.cpp b/torch/csrc/jit/script/init.cpp
@@ -233,6 +233,17 @@ static std::shared_ptr<Graph> _propagate_and_assign_input_and_output_shapes(
   return retval;
 }
 
+void addFunctionToModule(
+    Module& module,
+    const std::shared_ptr<Function>& func) {
+  // Make a graph with a fake self argument
+  auto graph = func->graph()->copy();
+  auto v = graph->insertInput(0, "self");
+  v->setType(module.module_object()->type());
+  module.module_object()->type()->compilation_unit().create_function(
+      "forward", graph);
+}
+
 void initJitScriptBindings(PyObject* module) {
   auto m = py::handle(module).cast<py::module>();
 
@@ -481,6 +492,27 @@ void initJitScriptBindings(PyObject* module) {
             }
             return result;
           })
+      .def(
+          "save",
+          [](std::shared_ptr<Function> self,
+             const std::string& filename,
+             const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
+            Module module;
+            addFunctionToModule(module, self);
+            module.save(filename, _extra_files);
+          },
+          py::arg("filename"),
+          py::arg("_extra_files") = ExtraFilesMap())
+      .def(
+          "save_to_buffer",
+          [](std::shared_ptr<Function> self,
+             const ExtraFilesMap& _extra_files = ExtraFilesMap()) {
+            std::ostringstream buf;
+            Module module;
+            addFunctionToModule(module, self);
+            return py::bytes(buf.str());
+          },
+          py::arg("_extra_files") = ExtraFilesMap())
       .def_property_readonly("graph", &Function::graph)
       .def_property_readonly("schema", &Function::getSchema)
       .def_property_readonly(
diff --git a/torch/csrc/jit/script/module.cpp b/torch/csrc/jit/script/module.cpp
@@ -195,7 +195,9 @@ std::pair<std::shared_ptr<Graph>, std::vector<Slot>> lower_graph(
     }
     if (e.n->kind() != prim::GetAttr) {
       throw ErrorReport(e.n->sourceRange())
-          << "temporary: the only valid use of a module is looking up an attribute";
+          << "temporary: the only valid use of a module is looking up an "
+             "attribute but found "
+          << *e.n;
     }
     Slot slot(e.mod, e.mod->type()->getAttributeSlot(e.n->s(attr::name)));
     if (ClassTypePtr c = e.n->output()->type()->cast<ClassType>()) {

Original file line number	Diff line number	Diff line change
`@@ -195,7 +195,9 @@ std::pair<std::shared_ptr<Graph>, std::vector<Slot>> lower_graph(`
`195`	`195`	`}`
`196`	`196`	`if (e.n->kind() != prim::GetAttr) {`
`197`	`197`	`throw ErrorReport(e.n->sourceRange())`
`198`		`- << "temporary: the only valid use of a module is looking up an attribute";`
	`198`	`+ << "temporary: the only valid use of a module is looking up an "`
	`199`	`+ "attribute but found "`
	`200`	`+ << *e.n;`
`199`	`201`	`}`
`200`	`202`	`Slot slot(e.mod, e.mod->type()->getAttributeSlot(e.n->s(attr::name)));`
`201`	`203`	`if (ClassTypePtr c = e.n->output()->type()->cast<ClassType>()) {`