Add section about .code to docs

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/18493

Differential Revision: D14634677

Pulled By: jamesr66a

fbshipit-source-id: 9ee065f6ce4218f725b93deb4c64b4ef55926145
This commit is contained in:
James Reed
2019-03-26 20:47:23 -07:00
committed by Facebook Github Bot
parent 45ec4920e3
commit f447b63ed0
2 changed files with 60 additions and 6 deletions

View File

@ -694,8 +694,56 @@ Disable JIT for Debugging
function.
Inspecting Code
^^^^^^^^^^^^^^^
TorchScript provides a code pretty-printer for all ScriptModule instances. This
pretty-printer gives an interpretation of the script method's code as valid
Python syntax. For example::
@torch.jit.script
def foo(len):
# type: (int) -> torch.Tensor
rv = torch.zeros(3, 4)
for i in range(len):
if i < 10:
rv = rv - 1.0
else:
rv = rv + 1.0
return rv
print(foo.code)
A ``ScriptModule`` with a single ``forward`` method will have an attribute
``code``, which you can use to inspect the ``ScriptModule``'s code.
If the ScriptModule has more than one method, you will need to access
``.code`` on the method itself and not the module. We can inspect the
code of a method named ``bar`` on a ScriptModule by accessing ``.bar.code``.
The example script abouve produces the code::
def forward(self,
len: int) -> Tensor:
rv = torch.zeros([3, 4], dtype=None, layout=None, device=None)
rv0 = rv
for i in range(len):
if torch.lt(i, 10):
rv1 = torch.sub(rv0, 1., 1)
else:
rv1 = torch.add(rv0, 1., 1)
rv0 = rv1
return rv0
This is TorchScript's interpretation of the code for the ``forward`` method.
You can use this to ensure TorchScript (tracing or scripting) has captured
your model code correctly.
Interpreting Graphs
^^^^^^^^^^^^^^^^^^^
TorchScript also has a representation at a lower level than the code pretty-
printer, in the form of IR graphs.
TorchScript uses a static single assignment (SSA) intermediate representation
(IR) to represent computation. The instructions in this format consist of
ATen (the C++ backend of PyTorch) operators and other primitive operators,
@ -714,11 +762,8 @@ Interpreting Graphs
print(foo.graph)
A ``ScriptModule`` with a single ``forward`` method will have an attribute
``graph``, which you can use to inspect the IR representing the computation.
If the ScriptModule has more than one method, you will need to access
``.graph`` on the method itself and not the module. We can inspect the
graph of a method named ``bar`` on a ScriptModule by accessing ``.bar.graph``.
``.graph`` follows the same rules described in the Inspecting Code section
with regard to ``forward`` method lookup.
The example script above produces the graph::

View File

@ -987,7 +987,16 @@ void initJitScriptBindings(PyObject* module) {
std::vector<ClassTypePtr> classes;
PythonPrint(oss, m, constants, classes, true);
return std::make_pair(oss.str(), std::move(constants));
});
})
.def_property_readonly(
"code",
[](Method& self) {
std::ostringstream ss;
std::vector<at::Tensor> tensors;
std::vector<ClassTypePtr> classes;
PythonPrint(ss, self, tensors, classes, false);
return ss.str();
});
m.def(
"_jit_script_compile",