attorch

attorch 项目介绍

attorch 是一个基于 PyTorch 的子集，完全使用 OpenAI 的 Triton 开发。该项目旨在创建一个易于修改和阅读的神经网络模块集合，同时在效率上与 PyTorch 保持一致或有所提高。通过简单直观的设计，attorch 为那些缺乏编写 CUDA 内核技术专长或资源的开发者提供了一个出色的起点。

项目目标

现有的 Triton 支持的 PyTorch 框架多聚焦于 NLP 应用，而 attorch 则希望能覆盖更多领域，特别是在计算机视觉方面。该项目支持前向和反向传播，因此可以用于模型训练和推理，但推理性能不如专门的推理引擎。

安装指南

要使用 attorch，需要安装 torch==2.4.0 和 triton==3.0.0 这两个依赖库，并克隆此项目的代码库开始使用。

图层功能

attorch 提供了一系列层功能，这些层与 PyTorch 中的等价层行为一致。以下是目前已实现的层，并支持自动混合精度（AMP）：

• attorch.Conv1d 和 attorch.Conv2d：分别用于一维和二维卷积操作，并可选择性地添加偏置。
• attorch.MultiheadAttention：应用多头自注意力机制。
• 各种激活函数如 Hardsigmoid、ReLU、LeakyReLU 等：对输入数据执行特定的非线性变换，可选择性融合 dropout。
• 正则化层如 BatchNorm1d 和 LayerNorm：进行批量或层归一化。
• 损失函数如 L1Loss 和 CrossEntropyLoss：用于计算输入与目标之间的误差。

数学函数

attorch 包含一些纯粹的数学函数，这些函数由 attorch.math 提供，旨在简化自定义内核和操作融合的实现。尽管这些函数只实现了前向传递，但由于其纯净性和无 I/O 操作特性，可以通过 triton-autodiff 库自动求导。

与 PyTorch 的兼容性

为了便于 attorch 与 PyTorch 层的结合，提供了 attorch.nn 接口。如果所需的层在 attorch 中不可用，将自动回退使用 PyTorch 的实现。

例如：

from attorch import nn

lin = nn.Linear(10, 20) # 使用 attorch 的线性层
gap = nn.AdaptiveAvgPool2d(1) # 由于 global pooling 在 attorch 中不可用，使用 PyTorch 的实现

测试

attorch 的每个模块都可以与其 PyTorch 的对应模块进行测试，以确保正确性。这些测试位于 tests/ 目录下，并可以使用 pytest 执行进行验证。虽然可能会由于数值精度问题导致某些测试失败，但在大多数实际使用案例中这并不是问题。

attorch 的目标是为开发者提供一个灵活且高效的平台，用于开发自定义的深度学习操作，帮助他们在学习和实现过程中获得更为便捷和直观的体验。这个项目既是现有技术的延续，又是对未来创新的激发。