attorch

• 简介
• 安装
• 层
• 数学函数
• PyTorch 回退
• 测试
• 示例
• 引用

简介

attorch 是 PyTorch 的 nn 模块的一个子集，完全用 Python 编写，并基于 OpenAI 的 Triton 实现。其目标是提供一个易于修改、自包含且可读性强的神经网络模块集合，同时保持或提升与 PyTorch 相当的效率。换句话说，它旨在成为一个可分支的项目，具有简单直观的设计，为那些希望开发自定义深度学习操作、但对纯 PyTorch 实现的速度不满意，且缺乏编写 CUDA 内核的技术专长或资源的人们，提供一个易于上手的起点。

目前已有许多优秀的基于 Triton 的类 PyTorch 框架，包括 kernl、xFormers、Unsloth 和 fla，但它们大多专注于 Transformer 和自然语言处理应用。相比之下，attorch 则更加全面，不仅涵盖 NLP 领域，还提供了计算机视觉等其他领域的多种层。此外，attorch 并非仅用于推理的包，它完全支持前向和反向传播，因此既可用于训练也可用于推理，尽管其在推理方面的性能通常不如专门的推理引擎。

安装

attorch 的唯一依赖是 torch==2.4.0 和 triton==3.0.0。请安装这两个库的指定版本，并克隆本仓库以开始使用。

层

目前已实现的层，并支持自动混合精度（AMP），包括：

• attorch.Conv1d: 使用权重对输入进行 1D 卷积，可选地添加偏置。
• attorch.Conv2d: 使用权重对输入进行 2D 卷积，可选地添加偏置。
• attorch.AvgPool1d: 对输入中的 1D 像素窗口取平均。
• attorch.AvgPool2d: 对输入中的 2D 像素窗口取平均。
• attorch.MultiheadAttention: 对输入应用多头缩放点积注意力机制。
• attorch.ELU: 对输入应用 ELU 激活函数，可选地融合 Dropout。
• attorch.Hardshrink: 对输入应用硬收缩激活函数，可选地融合 Dropout。
• attorch.Hardsigmoid: 对输入应用硬 Sigmoid 激活函数，可选地融合 Dropout。
• attorch.Hardswish: 对输入应用硬 Swish 激活函数，可选地融合 Dropout。
• attorch.Hardtanh: 对输入应用硬双曲正切激活函数，可选地融合 Dropout。
• attorch.LeakyReLU: 对输入应用 Leaky ReLU 激活函数，可选地融合 Dropout。
• attorch.LogSigmoid: 对输入应用 Sigmoid 函数的对数，可选地融合 Dropout。
• attorch.CELU: 对输入应用 CELU 激活函数，可选地融合 Dropout。
• attorch.GELU: 对输入应用 GELU 激活函数，可选地融合 Dropout。
• attorch.ReLU: 对输入应用 ReLU 激活函数，可选地融合 Dropout。
• attorch.ReLU6: 对输入应用 ReLU6 激活函数，可选地融合 Dropout。
• attorch.SELU: 对输入应用 SELU 激活函数，可选地融合 Dropout。
• attorch.SiLU: 对输入应用 SiLU 激活函数，可选地融合 Dropout。
• attorch.Mish: 对输入应用 Mish 激活函数，可选地融合 Dropout。
• attorch.Softplus: 对输入应用 Softplus 激活函数，可选地融合 Dropout。
• attorch.Softshrink: 对输入应用 Softshrink 激活函数，可选地融合 Dropout。
• attorch.Softsign: 对输入应用 Softsign 激活函数，可选地融合 Dropout。
• attorch.Sigmoid: 对输入应用 Sigmoid 激活函数，可选地融合 Dropout。
• attorch.Tanh: 对输入应用 Tanh 激活函数，可选地融合 Dropout。
• attorch.Tanhshrink: 对输入应用 Tanh 激活函数，可选地融合 Dropout。
• attorch.GLU: 对输入应用门控线性单元，可选地指定激活函数。
• attorch.LogSoftmax: 使用 Softmax 对输入进行归一化，并取其对数。
• attorch.Softmax: 使用 Softmax 对输入进行归一化。
• attorch.Softmin: 使用 Softmin 对输入进行归一化。
• attorch.BatchNorm1d: 对 2D 或 3D 输入进行批归一化，可选地融合激活函数，并在归一化前的结果上添加残差。
• attorch.BatchNorm2d: 对 4D 输入进行批归一化，可选地融合激活函数，并在归一化前的结果上添加残差。
• attorch.LayerNorm: 对输入进行层归一化。
• attorch.RMSNorm: 对输入进行均方根归一化。
• attorch.Linear: 使用权重对输入进行线性变换，可选地添加偏置并融合激活函数。
• attorch.Dropout: 在训练过程中随机将输入中的元素置零。
• attorch.L1Loss: 计算输入与目标之间的 L1 误差（平均绝对误差）。
• attorch.MSELoss: 计算输入与目标之间的平方 L2 误差（均方误差）。
• attorch.CrossEntropyLoss: 计算输入与目标之间的平均交叉熵损失，可选地对各类别进行加权。
• attorch.NLLLoss: 计算输入与目标之间的负对数似然损失，可选地对各类别进行加权。
• attorch.HuberLoss: 计算输入与目标之间的 Huber 损失。
• attorch.SmoothL1Loss: 计算输入与目标之间的平滑 L1 误差。

除非在其文档字符串中另有说明，上述各层的行为与其对应的 PyTorch 版本完全一致。

数学函数

Triton 内核通常由两部分组成：一部分负责相关张量的加载和存储，另一部分则使用适当的数学函数对数据进行变换。例如，层归一化内核会从输入中读取一行或多行数据（加载），对特征进行标准化处理（数学运算），然后将结果写入目标容器（存储）。attorch.math 提供了一系列纯数学函数，旨在简化自定义内核的实现以及操作融合。尽管 attorch.math 中目前仅实现了这些函数的前向传播，但由于其纯函数特性且不涉及 I/O 操作，其梯度可以通过 triton-autodiff 库自动推导出来。attorch 的许多内核都可以通过用对应的 attorch.math 变换或其导数替换其中的数学部分来重构，但这样做会牺牲 attorch 单文件、自包含的设计理念，因此 attorch.math 与库的其他部分仍将保持分离。

PyTorch 回退机制

为了更方便地将 attorch 与 PyTorch 层集成，提供了 attorch.nn 模块，它为 attorch 的模块提供了一个接口，并在所需层未实现时回退到 PyTorch 实现，如下所示：

from attorch import nn


lin = nn.Linear(10, 20) # 使用 attorch 的线性层
gap = nn.AdaptiveAvgPool2d(1) # 由于 attorch 中没有 GAP 层，因此使用 PyTorch 的全局池化层

尽管 attorch 自带卷积和池化层，但这些模块的性能相比 PyTorch 极为缓慢。因此，attorch.nn 暴露的是 PyTorch 的卷积和池化操作，而非 attorch 自身的实现。

测试

每个模块都可以与其对应的 PyTorch 实现进行对比测试，以确保正确性。这些测试位于 tests/ 目录下，可以使用 pytest 运行。此外，还提供了一个 --subset 选项，用于在较小的数据形状子集上运行测试，从而加快评估速度，但准确性会有所降低。需要注意的是，由于数值精度问题，部分测试可能会失败，但在大多数实际应用场景中，这通常不会构成问题。

示例

examples/ 目录下包含了一些常见的深度学习工作流，分别用 attorch 和 PyTorch 实现。虽然 attorch 可以作为 torch.nn 的直接替代品使用，但某些模块支持可选的内核融合功能，因此需要额外的参数，而这些参数与 PyTorch 对应的实现并不兼容。这些差异在示例中得到了展示。

引用

@inproceedings{tillet2019triton,
  title={Triton: an intermediate language and compiler for tiled neural network computations},
  author={Tillet, Philippe and Kung, Hsiang-Tsung and Cox, David},
  booktitle={Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages},
  pages={10--19},
  year={2019}
}

@Misc{xFormers2022,
  author =       {Benjamin Lefaudeux and Francisco Massa and Diana Liskovich and Wenhan Xiong and Vittorio Caggiano and Sean Naren and Min Xu and Jieru Hu and Marta Tintore and Susan Zhang and Patrick Labatut and Daniel Haziza and Luca Wehrstedt and Jeremy Reizenstein and Grigory Sizov},
  title =        {xFormers: A modular and hackable Transformer modelling library},
  howpublished = {\url{https://github.com/facebookresearch/xformers}},
  year =         {2022}
}

@software{yang2024fla,
  title  = {FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism},
  author = {Yang, Songlin and Zhang, Yu},
  url    = {https://github.com/fla-org/flash-linear-attention},
  month  = jan,
  year   = {2024}
}

attorch 快速上手指南

attorch 是一个基于 OpenAI Triton 纯 Python 实现的 PyTorch nn 模块子集。它旨在提供一套易读、可修改且自包含的神经网络组件，在保持或超越 PyTorch 效率的同时，为开发者提供自定义深度学习算子的便捷起点，无需编写复杂的 CUDA 代码。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：支持 Linux (推荐) 或 macOS。
• Python 版本：建议 Python 3.8+。
• 核心依赖：
  • torch == 2.4.0
  • triton == 3.0.0

注意：attorch 对版本有严格要求，请务必安装指定版本的 PyTorch 和 Triton，否则可能导致兼容性问题。

安装步骤

1. 安装依赖 建议使用国内镜像源加速安装过程。请运行以下命令安装指定版本的依赖：

   pip install torch==2.4.0 triton==3.0.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：如果您的 GPU 驱动较新或需要特定 CUDA 版本的 torch，请参考 PyTorch 官网获取对应的安装命令，但需确保最终版本为 2.4.0)

2. 克隆项目 由于 attorch 目前主要通过源码使用，请克隆仓库到本地：

   git clone https://github.com/BobMcDear/attorch.git
   cd attorch
   

3. 验证安装 进入项目目录后，您可以直接导入模块进行测试：

   python -c "import attorch; print('attorch imported successfully')"
   

基本使用

attorch 的设计初衷是作为 torch.nn 的替代品。您可以直接从 attorch 导入层，或者使用 attorch.nn 接口（该接口会在 attorch 未实现某层时自动回退到 PyTorch 原生实现）。

示例 1：直接使用 attorch 层

以下示例展示了如何使用 attorch 实现的线性层 (Linear) 和激活函数 (ReLU)：

import torch
from attorch import Linear, ReLU

# 定义输入张量 (Batch Size=2, Input Dim=10)
x = torch.randn(2, 10).cuda()

# 实例化 attorch 层
linear_layer = Linear(10, 20)
activation = ReLU()

# 前向传播
out = activation(linear_layer(x))

print(out.shape) # 输出: torch.Size([2, 20])

示例 2：使用混合回退模式 (推荐)

使用 attorch.nn 可以无缝混合使用 attorch 优化层和 PyTorch 原生层。例如，attorch 目前不包含自适应池化层，会自动调用 PyTorch 实现：

from attorch import nn

# 使用 attorch 优化的 Linear 层
lin = nn.Linear(10, 20) 

# 使用 PyTorch 原生的 AdaptiveAvgPool2d (因为 attorch 暂未实现)
gap = nn.AdaptiveAvgPool2d(1) 

print(f"Linear type: {type(lin)}")
print(f"Pool type: {type(gap)}")

特别说明

• 卷积与池化性能：虽然 attorch 实现了 Conv1d, Conv2d, AvgPool 等层，但其性能目前远慢于 PyTorch 原生实现。在使用 attorch.nn 时，卷积和池化操作会自动指向 PyTorch 版本以保证速度。
• 训练支持：attorch 完整支持前向和反向传播，可用于模型训练。
• 算子融合：部分激活函数层（如 ReLU, GELU）支持可选的 Dropout 融合，使用时可查看具体文档或示例代码以利用此特性。