torch.jit.interface¶
- torch.jit.interface(obj)[源代码]¶
装饰以注释不同类型的类或模块。
此装饰器可用于定义一个接口,该接口可用于注释不同类型的类或模块。这可用于注释可能具有不同类型(但实现相同接口)的子模块或属性类,或可在运行时交换的子模块或属性类;或存储不同类型的模块或类的列表。
它有时用于实现“可调用对象” - 实现接口但实现方式不同且可以互换的函数或模块。
示例: .. testcode
import torch from typing import List @torch.jit.interface class InterfaceType: def run(self, x: torch.Tensor) -> torch.Tensor: pass # implements InterfaceType @torch.jit.script class Impl1: def run(self, x: torch.Tensor) -> torch.Tensor: return x.relu() class Impl2(torch.nn.Module): def __init__(self) -> None: super().__init__() self.val = torch.rand(()) @torch.jit.export def run(self, x: torch.Tensor) -> torch.Tensor: return x + self.val def user_fn(impls: List[InterfaceType], idx: int, val: torch.Tensor) -> torch.Tensor: return impls[idx].run(val) user_fn_jit = torch.jit.script(user_fn) impls = [Impl1(), torch.jit.script(Impl2())] val = torch.rand(4, 4) user_fn_jit(impls, 0, val) user_fn_jit(impls, 1, val)