Python PyTorch Attribute用法及代碼示例

本文簡要介紹python語言中 torch.jit.Attribute 的用法。

用法:

class torch.jit.Attribute(value, type)

參數：

• value-要分配給屬性的初始值。

• type-Python 類型

返回：

返回value

此方法是一個返回 value 的 pass-through 函數，主要用於向 TorchScript 編譯器指示左側表達式是類型為 type 的類實例屬性。請注意，torch.jit.Attribute 隻能在nn.Module 子類的__init__ 方法中使用。

雖然 TorchScript 可以推斷出大多數 Python 表達式的正確類型，但在某些情況下類型推斷可能會出錯，包括：

• 空容器，例如 [] 和 {} ，其中 TorchScript 假定為 Tensor 的容器

• 可選類型，如 Optional[T] 但分配了 T 類型的有效值，TorchScript 會假設它是類型 T 而不是 Optional[T]

在 Eager 模式下，它隻是一個 pass-through 函數，它返回 value 而沒有其他含義。

例子：

import torch
from typing import Dict

class AttributeModule(torch.nn.Module):
    def __init__(self):
        super(M, self).__init__()
        self.foo = torch.jit.Attribute(0.1, float)

        # we should be able to use self.foo as a float here
        assert 0.0 < self.foo

        self.names_ages = torch.jit.Attribute({}, Dict[str, int])
        self.names_ages["someone"] = 20
        assert isinstance(self.names_ages["someone"], int)

m = AttributeModule()
# m will contain two attributes
# 1. foo of type float
# 2. names_ages of type Dict[str, int]