torch.jit.annotate

torch.jit.annotate(the_type, the_value)

用于在 TorchScript 编译器中指定 the_value 的类型。

此方法是一个直通函数，返回 the_value，用于向 TorchScript 编译器提示 the_value 的类型。在 TorchScript 外部运行时，此函数不起作用（no-op）。

尽管 TorchScript 可以推断大多数 Python 表达式的正确类型，但在某些情况下，类型推断可能不准确，包括：

• 空容器，如 [] 和 {}，TorchScript 默认将其视为 Tensor 的容器

• 可选类型，如 Optional[T]，但实际赋值为类型 T 的有效值时，TorchScript 会假定其类型为 T 而非 Optional[T]

注意，annotate() 在 torch.nn.Module 子类的 __init__ 方法中不起作用，因为它是在 Eager 模式下执行的。要标注 torch.nn.Module 属性的类型，请改用 Attribute()。

示例

import torch
from typing import Dict

@torch.jit.script
def fn():
    # Telling TorchScript that this empty dictionary is a (str -> int) dictionary
    # instead of default dictionary type of (str -> Tensor).
    d = torch.jit.annotate(Dict[str, int], {})

    # Without `torch.jit.annotate` above, following statement would fail because of
    # type mismatch.
    d["name"] = 20

参数

• the_type – 应传递给 TorchScript 编译器作为 the_value 类型提示的 Python 类型

• the_value – 需要提示类型的 值 或 表达式。

返回值

the_value 作为返回值传回。