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())