torch.jit.interface

torch.jit.interface(obj)

用於裝飾以註解不同類型的類別或模組。

此裝飾器可用於定義一個介面，該介面可用於註解不同類型的類別或模組。 這可用於註解一個子模組或屬性類別，該子模組或屬性類別可能具有實現相同介面的不同類型，或者可以在運行時進行交換；或儲存不同類型的模組或類別的清單。

它有時用於實現“可呼叫物件 (Callables)” - 實現介面但其實作不同的函數或模組，並且可以交換。

範例： .. 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)