# attorch
• **[介绍](#介绍)**<br>
• **[安装](#安装)**<br >
• **[层](#层)**<br >
• **[数学函数](#数学函数)**<br >
• **[PyTorch 回退](#pytorch-回退)**<br >
• **[测试](#测试)**<br >
## 介绍
attorch 是 PyTorch 的 [```nn```](https://pytorch.org/docs/stable/nn.html) 模块的一个子集,纯粹使用 OpenAI 的 [Triton](https://github.com/openai/triton) 以 Python 编写。它的目标是成为一个易于修改、自包含且可读的神经网络模块集合,同时保持或提高 PyTorch 的效率。换句话说,它旨在成为一个具有简单、直观设计的可分叉项目,能够作为那些希望开发自定义深度学习操作但对纯 PyTorch 实现的速度不满意且没有技术专业知识或资源编写 CUDA 内核的人们的一个易于访问的起点。
已经存在许多由 Triton 支持的优秀的类似 PyTorch 的框架,包括 [kernl](https://github.com/ELS-RD/kernl/tree/main)、[xFormers](https://github.com/facebookresearch/xformers)、[Unsloth](https://github.com/unslothai/unsloth) 和 [```fla```](https://github.com/sustcsonglin/flash-linear-attention),但大多数主要集中在 Transformers 和 NLP 应用上,而 attorch 旨在更具包容性,也涵盖与 NLP 之外的领域有关的各种层,如计算机视觉。此外,attorch 不是一个仅用于推理的包,完全支持前向和后向传递,这意味着它可以在训练和推理过程中使用,尽管在推理方面它的性能通常不如专用推理引擎。
## 安装
attorch 的唯一依赖是 ```torch==2.2.0``` 和 ```triton==2.2.0```。请安装这两个库的指定版本并克隆此仓库以开始使用。
## 层
目前实现的层,具有自动混合精度(AMP)支持,分别是,
* ```attorch.Conv1d```: 使用权重对输入进行一维卷积,可选地添加偏置项。
* ```attorch.Conv2d```: 使用权重对输入进行二维卷积,可选地添加偏置项。
* ```attorch.MultiheadAttention```: 对输入应用多头缩放点积注意力。
* ```attorch.GELU```: 对输入应用 GELU,可选地融合 dropout。
* ```attorch.ReLU```: 对输入应用 ReLU,可选地融合 dropout。
* ```attorch.SiLU```: 对输入应用 SiLU,可选地融合 dropout。
* ```attorch.Sigmoid```: 对输入应用 Sigmoid,可选地融合 dropout。
* ```attorch.Tanh```: 对输入应用 Tanh,可选地融合 dropout。
* ```attorch.GLU```: 使用任意激活函数对输入应用门控线性单元(GLU)。
* ```attorch.LogSoftmax```: 使用 Softmax 对输入进行归一化并取其对数。
* ```attorch.Softmax```: 使用 Softmax 对输入进行归一化。
* ```attorch.Softmin```: 使用 Softmin 对输入进行归一化。
* ```attorch.BatchNorm1d```: 对二维或三维输入进行批量归一化,可选地融合激活函数并向预激活结果添加残差。
* ```attorch.BatchNorm2d```: 对四维输入进行批量归一化,可选地融合激活函数并向预激活结果添加残差。
* ```attorch.LayerNorm```: 对输入进行层归一化。
* ```attorch.Linear```: 使用权重对输入进行线性变换,可选地添加偏置并融合激活函数。
* ```attorch.Dropout```: 在训练期间随机将输入中的元素置零。
* ```attorch.L1Loss```: 测量输入和目标之间的平均绝对误差。
* ```attorch.MSELoss```: 测量输入和目标之间的均方误差。
* ```attorch.CrossEntropyLoss```: 测量输入和目标之间的平均交叉熵损失,可选地对每个类别重新加权。
* ```attorch.NLLLoss```: 测量输入和目标之间的负对数似然损失,可选地对每个类别重新加权。
除非在 docstring 中另有说明,否则上述层的行为与 PyTorch 等效层完全相同。
## 数学函数
Triton 内核通常由两个部分组成:一个处理相关张量的加载和存储,另一个使用适当的数学函 数转换数据。例如,一个层归一化内核读取输入的一行或几行(加载),对特征进行标准化(数学),并将结果写入容器(存储)。这些纯数学函数的一部分由 ```attorch.math``` 提供,目的是促进自定义内核和操作融合的实现。尽管上述函数的前向传递仅在 ```attorch.math``` 中可用,但由于它们的纯净性和没有 I/O 操作,可以通过 [```triton-autodiff```](https://github.com/srush/triton-autodiff) 库自动推导其梯度。值得注意的是,attorch 的内核的很大一部分可以通过用相应的 ```attorch.math``` 转换或其导数取代其数学部分进行重 构,但这样做会牺牲 attorch 的单文件和自包含设计,因此 ```attorch.math``` 和 attorch 的其余部分将保持分离。
## PyTorch 回退
为了更容易地集成 attorch 和 PyTorch 层,提供了 ```attorch.nn```,它提供了一个接口来使用 attorch 的模块并在需要的层不可用时回退到 PyTorch,如下所示。
```python
from attorch import nn
lin = nn.Linear(10, 20) # 使用 attorch 的线性层
gap = nn.AdaptiveAvgPool2d(1) # 使用 PyTorch 的全局池化,因为 attorch 中没有提供 GAP
测试
每个模块都可以与其 PyTorch 对应模块进行测试以确保正确性。这些测试包含在 tests/
目录下,并可以使用 pytest
执行。需要注意的是,由于数值精度问题,某些测试可能会失败,但在大多数实际用例中,这不应成为问题。