torchtyping

项目介绍: Torchtyping

Torchtyping 是一个用于 Tensor（张量）类型注释的工具，专注于描述张量的形状、数据类型（dtype）、名字等各种属性。它为开发者提供了一种更清晰且更安全的方式编写代码，帮助程序员避免常见的错误。

主要功能

在传统的编写代码过程中，开发者通常需要通过添加大量的注释，例如“# x has shape (batch, hidden_state)”或使用“assert x.shape == y.shape”这样的语句来追踪张量的形状。Torchtyping 通过类型注释简化了这一过程，使代码更加简洁，并减少了错误的发生。

利用 torchtyping，开发者可以如下定义函数：

def batch_outer_product(x:   TensorType["batch", "x_channels"],
                        y:   TensorType["batch", "y_channels"]
                        ) -> TensorType["batch", "x_channels", "y_channels"]:
    return x.unsqueeze(-1) * y.unsqueeze(-2)

安装方法

您可以通过以下命令来安装 Torchtyping:

pip install torchtyping

注意，需要 Python 3.7 以上版本，以及 PyTorch 1.7.0 及以上版本。如果您使用 typeguard，则其版本需要低于 3.0.0。

使用指南

Torchtyping 可以注释以下属性：

• 形状: 张量的维度和大小；
• 数据类型 (dtype): 例如浮点型、整数等；
• 布局 (layout): 稠密或稀疏布局；
• 维度的名字: 针对命名张量；
• 任意数量的批量维度通过使用 ...；
• 其他自定义属性。

在通过安装 typeguard 后，Torchtyping 可以在运行时验证这些类型，以确保张量与注释中的形状和数据类型一致。

示例代码

通过以下代码可了解如何在程序中应用：

from torch import rand
from torchtyping import TensorType, patch_typeguard
from typeguard import typechecked

patch_typeguard()  # 在 @typechecked 之前使用

@typechecked
def func(x: TensorType["batch"],
         y: TensorType["batch"]) -> TensorType["batch"]:
    return x + y

func(rand(3), rand(3))  # 正常运行
func(rand(3), rand(1))  # 报错：不同大小的 'batch' 维度

核心 API

Torchtyping 的核心是 TensorType，可以用来注释张量的多个属性组合：

• shape（形状）支持指定具体尺寸、名称或任意维度等；
• dtype（数据类型）可以是特定的 PyTorch 数据类型或者 Python 的基本数据类型；
• layout 说明张量是稠密还是稀疏的；
• details 用于传递自定义标志，如 is_named 或 is_float。

与 typeguard 的集成

通过 torchtyping.patch_typeguard() 函数，Torchtyping 可以与 typeguard 集成，实现实时类型检查。该函数应全局调用，确保 typeguard 能够适配 TensorType。

进一步文档

如果想深入了解如何编写 Torchtyping 自定义扩展、兼容 flake8 和 mypy 等，建议用户查阅 Torchtyping 的进一步文档。这些资源将为使用者提供更广泛的帮助，节省调试时间，并提高代码质量。

总结来说，Torchtyping 不仅使代码更加清晰易懂，而且通过正则化的类型注释提高了代码的鲁棒性，是每个使用 PyTorch 编写深度学习模型开发者的好帮手。