How can you script a python function?
foo_scripted = torch.jit.script(foo)
@torch.jit.script
def foo2(x, y):
if x.max() > y.max():
r = x
else: r = y
return r
# using a decorator
TorchScript excepts static typing, dynamic changes of variables are not allowed
If we want to script non-tensor objects, we have to include a type hint
How can you script a torch module?
my_module_scripted = torch.jit.script(MyModule(5))
# this will automatically script the forward() method
# and all methods invoked by it
How can you trace python code?
foo_traced = torch.jit.trace(foo, example_inputs=(tensor1, tensor2))
Tracing does not take care of control flow, and therefore only evaluates the code which was executed with the initial input
How can save, load and inspect TorchScript code
foo_traced.save("10_torch_script_module.pt")
loaded_script_module = torch.jit.load("10_torch_script_module.pt")
foo_scripted.graph
Zuletzt geändertvor 2 Jahren