torch.jit.ignore

torch.jit.ignore(drop=False, **kwargs)

这个装饰器向编译器指示一个函数或方法应该被忽略并保持为Python函数。这允许你在模型中保留尚未与TorchScript兼容的代码。如果在TorchScript中调用，被忽略的函数将把调用分派给Python解释器。带有被忽略函数的模型无法导出；请改用@torch.jit.unused。

示例（在方法上使用 @torch.jit.ignore）：

import torch
import torch.nn as nn

class MyModule(nn.Module):
    @torch.jit.ignore
    def debugger(self, x):
        import pdb
        pdb.set_trace()

    def forward(self, x):
        x += 10
        # 编译器通常会尝试编译 `debugger`，
        # 但由于它被 `@ignore` 了，它将被保留为一个
        # 对 Python 的调用
        self.debugger(x)
        return x

m = torch.jit.script(MyModule())

# 错误！调用 `debugger` 无法保存，因为它调用了 Python
m.save("m.pt")

示例（在方法上使用 @torch.jit.ignore(drop=True)）：

import torch
import torch.nn as nn

class MyModule(nn.Module):
    @torch.jit.ignore(drop=True)
    def training_method(self, x):
        import pdb
        pdb.set_trace()

    def forward(self, x):
        if self.training:
            self.training_method(x)
        return x

m = torch.jit.script(MyModule())

# 这是可以的，因为 `training_method` 没有被保存，调用被替换为 `raise`。
m.save("m.pt")