邻域注意力Transformer

基于滑动窗口注意力的强大分层视觉Transformer模型。

邻域注意力（NA，局部注意力）在我们的原始论文中被提出，即NAT，并可通过我们对PyTorch的扩展NATTEN高效运行。

我们最近又提出了一个新的模型，DiNAT，它通过扩张邻域来扩展NA（DiNA，稀疏全局注意力，又称扩张的局部注意力）。

NA/DiNA的组合能够保持局部性、维持平移等变性、以指数级方式扩大感受野，并捕捉更长距离的相互依赖关系，从而在下游视觉任务中带来显著的性能提升，例如用于图像生成的StyleNAT。

新闻

2023年3月25日

• 邻域注意力Transformer已被CVPR 2023接收！

2022年11月18日

• NAT和DiNAT现已通过HuggingFace的transformers库提供。
  • NAT和DiNAT的分类模型也在HuggingFace模型库中上线：NAT | DiNAT

2022年11月11日

• 新预印本：StyleNAT：为每个头赋予全新视角。
  • 基于邻域注意力的风格化GAN在FFHQ-256数据集上以2.05的FID值创下新的SOTA记录。

2022年10月8日

• NATTEN现已作为pip包发布！
  • 现在你可以通过预编译的wheel文件安装NATTEN，并在几秒钟内开始使用。
  • NATTEN将作为一个独立项目维护和发展，以支持更广泛的滑动窗口注意力应用，甚至超越计算机视觉领域。

2022年9月29日

• 新预印本：扩张邻域注意力Transformer。

张化邻域注意力 :fire:

一种基于邻域注意力（局部注意力）和扩张邻域注意力（稀疏全局注意力）的新分层视觉Transformer模型，在下游任务中表现出显著的性能提升。

请查看DiNAT README。

邻域注意力Transformer

我们的原始论文，邻域注意力Transformer (NAT)，是首个高效的滑动窗口局部注意力机制。

邻域注意力的工作原理

邻域注意力会将查询token（红色）的感受野限定在其键值对（绿色）中的最近邻近token范围内。当邻域大小与图像尺寸相同时，这等同于点积自注意力。请注意，边缘区域属于特殊情况。

引用

@inproceedings{hassani2023neighborhood,
	title        = {邻域注意力Transformer},
	author       = {阿里·哈萨尼和史蒂文·沃尔顿和李嘉辰和李申和希普里·施伊},
	booktitle    = {IEEE/CVF计算机视觉与模式识别会议（CVPR）论文集},
	month        = {6月},
	year         = {2023},
	pages        = {6185-6194}
}
@article{hassani2022dilated,
	title        = {扩张邻域注意力Transformer},
	author       = {阿里·哈萨尼和希普里·施伊},
	year         = 2022,
	url          = {https://arxiv.org/abs/2209.15001},
	eprint       = {2209.15001},
	archiveprefix = {arXiv},
	primaryclass = {cs.CV}
}
@article{walton2022stylenat,
	title        = {StyleNAT：为每个头提供新视角},
	author       = {史蒂文·沃尔顿和阿里·哈萨尼和徐兴谦和王张阳和希普里·施伊},
	year         = 2022,
	url          = {https://arxiv.org/abs/2211.05770},
	eprint       = {2211.05770},
	archiveprefix = {arXiv},
	primaryclass = {cs.CV}
}

Neighborhood Attention Transformer 快速上手指南

Neighborhood Attention Transformer (NAT) 是一种基于滑动窗口注意力机制的高效分层视觉 Transformer。它通过局部注意力（Neighborhood Attention）和扩张局部注意力（Dilated Neighborhood Attention）在保持平移等变性的同时，有效扩大感受野，显著提升图像分类、分割及生成等下游任务的性能。

环境准备

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

• 操作系统: Linux (推荐 Ubuntu 18.04+) 或 macOS。Windows 支持取决于 NATTEN 的编译环境。
• Python: 3.8 或更高版本。
• PyTorch: 1.9 或更高版本（建议安装与您的 CUDA 版本匹配的最新版）。
• CUDA: 如果使用 GPU 加速，需安装对应的 CUDA Toolkit 和 cuDNN。
• 编译器: 需要安装 g++ 和 nvcc (若使用 GPU) 以编译 CUDA 扩展。

前置依赖检查： 确保已安装基础深度学习库：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 请将 cu118 替换为您实际的 CUDA 版本，如 cu121, cpu 等

安装步骤

NAT 的核心算子依赖于 NATTEN (Neighborhood Attention Extension)。为了获得最佳性能，强烈建议通过预编译的 wheel 包进行安装。

1. 安装 NATTEN

您可以直接从 PyPI 安装预编译版本（最快方式）：

pip install natten

注：如果上述命令失败或需要特定版本，请访问 NATTEN GitHub 获取针对特定 CUDA 版本的安装指令。国内用户若遇到下载慢的问题，可尝试使用清华源或阿里源：

pip install natten -i https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 Transformers 集成 (可选但推荐)

自 2022 年 11 月起，NAT 和 DiNAT 模型已集成到 Hugging Face transformers 库中。如果您希望通过 HF 接口调用模型：

pip install transformers

国内加速方案：

pip install transformers -i https://pypi.tuna.tsinghua.edu.cn/simple
export HF_ENDPOINT=https://hf-mirror.com

3. 克隆源码 (用于训练或研究底层代码)

如果您需要运行官方提供的训练脚本或研究具体实现：

git clone https://github.com/SHI-Labs/Neighborhood-Attention-Transformer.git
cd Neighborhood-Attention-Transformer
pip install -r requirements.txt

基本使用

方式一：通过 Hugging Face Transformers 加载预训练模型

这是最简单的推理方式，适用于图像分类任务。

from transformers import AutoImageProcessor, AutoModelForImageClassification
from PIL import Image
import torch

# 加载处理器和模型 (以 NAT-Tiny 为例)
model_name = "shi-labs/nat-tiny-224" # 可在 HuggingFace Hub 查找更多变体：nat-small, nat-base, dinat 等
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# 准备图像
image = Image.open("your_image.jpg").convert("RGB")

# 预处理并推理
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

# 获取预测结果
predicted_class_idx = logits.argmax(-1).item()
print(f"Predicted class: {model.config.id2label[predicted_class_idx]}")

方式二：直接使用 NATTEN 构建自定义层

如果您想在自定义模型架构中使用 Neighborhood Attention 算子：

import torch
from natten import NeighborhoodAttention2d

# 定义参数
dim = 64          # 嵌入维度
kernel_size = 7   # 邻域大小 (k)
dilation = 1      # 膨胀率 (1 为普通局部注意力，>1 为扩张注意力)
num_heads = 8     # 注意力头数

# 初始化注意力层
na_layer = NeighborhoodAttention2d(
    dim=dim,
    kernel_size=kernel_size,
    dilation=dilation,
    num_heads=num_heads
)

# 模拟输入张量 (Batch, Channels, Height, Width)
# 注意：NATTEN 通常期望输入格式为 (B, H, W, C) 或在内部处理转换，具体请参考 natten 文档
# 此处演示标准卷积风格输入，实际使用时请根据 natten 最新文档确认输入维度要求
# 对于 v0.14+，输入通常为 (B, H, W, C)
x = torch.randn(2, 32, 32, dim) 

# 前向传播
output = na_layer(x)

print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")

方式三：加载 DiNAT (扩张邻域注意力) 模型

DiNAT 是 NAT 的进阶版，结合了稀疏全局注意力，适合更复杂的场景。

from transformers import AutoModelForImageClassification

# 加载 DiNAT 模型
model_name = "shi-labs/dinat-tiny-224"
model = AutoModelForImageClassification.from_pretrained(model_name)

print(f"Loaded model: {model_name}")
# 后续推理步骤与上述 NAT 示例相同