mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
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:
committed by
Facebook Github Bot
parent
45ec4920e3
commit
f447b63ed0
@ -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::
|
||||
|
||||
|
@ -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",
|
||||
|
Reference in New Issue
Block a user