Python PyTorch export用法及代码示例

本文简要介绍python语言中 torch.jit.export 的用法。

用法:

torch.jit.export(fn)

此装饰器指示nn.Module 上的方法用作 ScriptModule 的入口点，应编译。

forward 隐式假定为入口点，因此不需要此装饰器。从forward 调用的函数和方法按照编译器看到的方式编译，因此它们也不需要这个装饰器。

示例(在方法上使用@torch.jit.export)：

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def implicitly_compiled_method(self, x):
        return x + 99

    # `forward` is implicitly decorated with `@torch.jit.export`,
    # so adding it here would have no effect
    def forward(self, x):
        return x + 10

    @torch.jit.export
    def another_forward(self, x):
        # When the compiler sees this call, it will compile
        # `implicitly_compiled_method`
        return self.implicitly_compiled_method(x)

    def unused_method(self, x):
        return x - 20

# `m` will contain compiled methods:
#     `forward`
#     `another_forward`
#     `implicitly_compiled_method`
# `unused_method` will not be compiled since it was not called from
# any compiled methods and wasn't decorated with `@torch.jit.export`
m = torch.jit.script(MyModule())