深入解析 Transformers 中的 Dilated Neighborhood Attention Transformer（DiNAT）架构与实战

导读

DiNAT（Dilated Neighborhood Attention Transformer）是 2022 年由 SHI-Labs 提出的分层视觉 Transformer 模型，其核心创新是引入 Dilated Neighborhood Attention（DiNA）——一种对 Neighborhood Attention（NA）的稀疏膨胀扩展，可在不增加计算成本的前提下指数级扩大感受野、捕获更全局的上下文。本文以 Hugging Face Transformers 当前仓库中的 DiNAT 官方文档 为主线，结合 src/transformers/models/dinat/ 下的配置与建模源码，系统讲解 DiNAT 的设计动机、NATTEN 依赖安装、推理与微调实战、完整配置项说明，以及将 DiNAT 作为骨干网络（backbone）接入检测、分割等框架的方法。

DiNAT 是什么：从 Neighborhood Attention 到 Dilated Neighborhood Attention

背景：局部注意力与自注意力的两难

视觉 Transformer 在降低自注意力二次方复杂度时，普遍采用局部注意力机制，例如滑窗式的 Neighborhood Attention（NA）和 Swin Transformer 的 Shifted Window Self-Attention。但论文指出，局部化注意力同时削弱了自注意力最有价值的两个特性：长程依赖建模能力与全局感受野。

DiNAT 的核心思想

DiNAT 由 Ali Hassani 与 Humphrey Shi 在论文 Dilated Neighborhood Attention Transformer 中提出（原文档标注该模型于 2022-09-29 发表论文、2022-11-18 由作者贡献进 Transformers）。它把 Neighborhood Attention 扩展为 Dilated Neighborhood Attention（DiNA）：让每个 token 关注其邻域内每隔若干像素采样一次的 token，即用膨胀率（dilation）拉开采样间距。这样一来：

• NA 的局部细粒度注意力负责捕捉局部细节；
• DiNA 的稀疏全局注意力负责在相隔更远的 token 间建立关联；
• 两者天然互补，膨胀后的感受野随层数指数级扩大，而注意力计算量保持不变（因为采样点数量未增加）。

从仓库源码中可以看出这种"零额外成本"的落地点：邻域注意力实现 将膨胀率作为参数直接传给 NATTEN 底层算子 natten2dqkrpb（QK 注意力分数）与 natten2dav（注意力加权），kernel 形状、参与计算的 token 数量都受 kernel_size 约束，与膨胀率无关：

attention_scores = natten2dqkrpb(query_layer, key_layer, self.rpb, self.kernel_size, self.dilation)
attention_probs = nn.functional.softmax(attention_scores, dim=-1)
attention_probs = self.dropout(attention_probs)
context_layer = natten2dav(attention_probs, value_layer, self.kernel_size, self.dilation)

论文中报告的主要结果（转述自原文档引用摘要）

论文报告 DiNAT 在多个强基线上取得明显改进，其 large 变体在 COCO 目标检测上比对应的 Swin large 快且高 1.5% box AP、在 COCO 实例分割上高 1.3% mask AP、在 ADE20K 语义分割上高 1.1% mIoU；并称在配以新框架后在 COCO（58.2 PQ）、ADE20K（48.5 PQ）上刷新全景分割 SOTA，在 Cityscapes（44.5 AP）与 ADE20K（35.4 AP）上成为实例分割 SOTA（均未使用额外数据）。需要说明的是，以上数值属于原论文的自述结论，引用时应以论文原文为准。

分层架构总览：四个阶段、逐级膨胀

与 NAT、Swin 类似，DiNAT 是标准的分层（hierarchical）视觉 Transformer。从 DinatEncoder 的实现看，模型依次经过 stem 与四个 stage，feature map 分辨率逐级减半、通道数逐级翻倍，各 stage 通过卷积下采样层衔接。

Patch Embedding（stem）

DinatPatchEmbeddings 用两层 stride=2 的 3×3 卷积实现 4×4 patch 化输入图像，将 (batch, num_channels, height, width) 转为 (batch, height/4, width/4, embed_dim)。该代码在配置的 patch_size 不等于 4 时会直接抛出 ValueError——这正是原文档"当前仅支持 patch size 4"提示的源码出处。

层与阶段的组织方式

每个 stage 内部按 depths 指定的层数堆叠 DinatLayer，每层使用一个指定的膨胀率；不同层之间膨胀率交替变化，形成类似 1, 8, 1 或 1, 4, 1, 4 的排布。相邻 stage 之间由 DinatDownsampler 完成 2 倍分辨率下采样：Conv2d(dim, 2*dim, 3×3, stride=2) 后接 LayerNorm。

DinatLayer 是一个标准的 pre-LayerNorm Transformer Block：LayerNorm → NeighborhoodAttention → DropPath → LayerNorm → MLP → DropPath 两条残差路径，并支持可选的 LayerScale（当 layer_scale_init_value > 0 时启用）。每个注意力层还会依据"膨胀窗口尺寸 = kernel_size × dilation"对空间不足的 feature map 做 padding（maybe_pad），并在计算后裁剪回原始尺寸。

默认配置下四个 stage 的膨胀率序列为（见 DinatConfig 默认值）：

Stage	层数（depths）	膨胀率序列（dilations）	膨胀窗口跨度
stage 1	3	[1, 8, 1]	7 / 56 / 7
stage 2	4	[1, 4, 1, 4]	7 / 28 / 7 / 28
stage 3	6	[1, 2, 1, 2, 1, 2]	7 / 14 / 7 / 14 / 7 / 14
stage 4	5	[1, 1, 1, 1, 1]	7 / 7 / 7 / 7 / 7

其中"窗口跨度"按默认 kernel_size=7 计算（窗口尺寸 = kernel×dilation）。可见设计刻意让浅层（高分辨率、token 密集）使用大膨胀率以捕捉长程关系，深层（低分辨率、感受野已足够大）回归普通邻域注意力。

环境准备：安装 NATTEN 内核库

原文档与源码都明确指出：DiNAT 无法脱离 NATTEN 运行，因为其前向计算依赖 NATTEN 对 NA/DiNA 的高效 CUDA 实现。

• 在 DinatModel 与 DinatBackbone 的构造器中都调用了 requires_backends(self, ["natten"])，环境缺库时会直接抛错；
• NATTEN 缺失提示 给出的安装方式如下：

# 方式一：使用预编译 wheel（Linux，推荐，见 shi-labs.com/natten 的发行说明）
# 方式二：源码编译（耗时较长）
pip install natten

安装约束与注意事项：

• Linux 优先：预编译 wheel 面向 Linux 提供；从源码编译也主要在 Linux 上验证。
• Windows 暂不支持：NATTEN 目前不提供 Windows 支持，这是引入 DiNAT 前需要评估的硬性前提。
• 安装完成后若在运行期导入仍报错，通常需要重启运行时（runtime）。

快速上手：用 DiNAT 做图像分类推理

DiNAT 可直接以官方预训练权重加载。配置文档标注的 mini 规格 checkpoint 为 shi-labs/dinat-mini-in1k-224（该命名空间下还发布有 small/base/large 等规格），完整推理流程如下：

from transformers import AutoImageProcessor, DinatForImageClassification
from PIL import Image

processor = AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224")
model = DinatForImageClassification.from_pretrained("shi-labs/dinat-mini-in1k-224")

# 读取本地图片
image = Image.open("cat.jpg")

inputs = processor(image, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits

predicted_class_idx = logits.argmax(-1).item()
print(model.config.id2label[predicted_class_idx])

原文档的 Resources 一节表明，DinatForImageClassification 由 examples/pytorch/image-classification 下的示例脚本与配套 notebook 支持；完整的图像分类数据处理、训练与评估流程，可继续参考 图像分类任务指南。

作为骨干网络（Backbone）使用

原文档明确 DiNAT 的典型定位是 backbone，即作为特征提取器为 DETR、MaskFormer 等下游视觉框架提供多尺度特征图。仓库为此专门提供了 DinatBackbone（继承 BackboneMixin），其 docstring 用法示例如下（此处以官方 NAT checkpoint 的处理器示意流程，加载 DiNAT 权重同理）：

from transformers import AutoImageProcessor, AutoBackbone

processor = AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224")
model = AutoBackbone.from_pretrained(
    "shi-labs/nat-mini-in1k-224",
    out_features=["stage1", "stage2", "stage3", "stage4"],
)

inputs = processor(image, return_tensors="pt")
outputs = model(**inputs)
feature_maps = outputs.feature_maps
print(list(feature_maps[-1].shape))  # 例如 [1, 512, 7, 7]

从实现看，DinatBackbone 会为每个被选中的 stage 额外套一个 LayerNorm（self.hidden_states_norms），并通过 out_features / out_indices 控制输出哪些 stage 的特征图（这两个字段也可在 DinatConfig 构造时传入）。最终 BackboneOutput.feature_maps 中的每个元素形状为 (batch, num_channels, height, width)。

模型输出解读：hidden_states 与 reshaped_hidden_states

原文档给出的一个高频注意点是输出张量的维度布局：

当 output_hidden_states=True 时，模型会同时输出 hidden_states 和 reshaped_hidden_states。其中 reshaped_hidden_states 的形状是 (batch, num_channels, height, width)，而非 (batch, height, width, num_channels)。

对应的机制在 DinatEncoder.forward 中非常直白：hidden_states 以 channels-last 布局（b, h, w, c）累积，同时用一行 permute(0, 3, 1, 2) 生成 channels-first 的副本存入 reshaped_hidden_states。

各输出容器的含义（以 DinatModelOutput 为准）：

字段	形状	含义
last_hidden_state	(batch, h, w, channels)	最后一层经过 LayerNorm 的隐藏状态
pooler_output	(batch, hidden_size)	最后一层隐藏状态的平均池化（add_pooling_layer=True 时返回）
hidden_states	每阶段各一份	每个 stage 输出（含 embedding 输出）
reshaped_hidden_states	每阶段各一份	同 hidden_states 但转置为 (b, c, h, w)
attentions	每层各一份	各层注意力权重（output_attentions=True 时返回）

DinatForImageClassification 的分类头即取 outputs[1]（池化后的向量）过一层 Linear，见 分类模型 forward；当传入 labels 时按 config.num_labels 自动计算交叉熵（多类）或 MSE（num_labels==1 回归）。

DinatConfig 完整参数解析

配置文件位于 configuration_dinat.py，默认值即对应 shi-labs/dinat-mini-in1k-224 规格。下表覆盖 DinatConfig 的公开字段：

参数	默认值	说明
patch_size	4	Patch 大小，当前实现仅支持 4（stem 由两层 stride-2 卷积构成）
num_channels	3	输入图像通道数（RGB）
embed_dim	64	Stem 输出通道数，也是 stage 1 的维度
depths	(3, 4, 6, 5)	四个 stage 各自堆叠的层数
num_heads	(2, 4, 8, 16)	四个 stage 各自的注意力头数
kernel_size	7	邻域注意力 kernel 尺寸
dilations	[[1,8,1],[1,4,1,4],[1,2,1,2,1,2],[1,1,1,1,1]]	每个 stage 内各层的膨胀率，长度须与 depths 对应
mlp_ratio	3.0	MLP 隐藏层相对维度倍率
qkv_bias	True	Q/K/V 投影是否带偏置
hidden_dropout_prob	0.0	隐藏层 dropout 概率
attention_probs_dropout_prob	0.0	注意力权重 dropout 概率
drop_path_rate	0.1	随机深度（stochastic depth）概率，随层线性递增
hidden_act	"gelu"	MLP 激活函数
initializer_range	0.02	权重初始化标准差
layer_norm_eps	1e-5	LayerNorm epsilon
layer_scale_init_value	0.0	LayerScale 初始值，0 表示关闭
out_features / out_indices	None	Backbone 输出 stage 的筛选（名称或索引）

有几个字段是自动推导而非直接设置的：

• num_layers：构造时由 len(depths) 自动算出（见 __post_init__）；
• hidden_size：为兼容 VisionEncoderDecoderModel 而设置，等于 embed_dim × 2^(stages-1)，即最后一个 stage 的通道数；
• 映射别名：num_attention_heads→num_heads、num_hidden_layers→num_layers（attribute_map），方便与其他模型 API 对齐。

手动构建配置与随机初始化

from transformers import DinatConfig, DinatModel

# 使用默认配置（dinat-mini 风格）
configuration = DinatConfig()

# 依据配置创建随机初始化模型
model = DinatModel(configuration)

# 访问模型配置
configuration = model.config

微调示例：加载预训练权重继续训练

从 ImageNet 预训练权重出发进行下游微调，是 DiNAT 最常见的应用路径。以自定义 10 类数据集为例：

from transformers import (
    AutoImageProcessor,
    DinatForImageClassification,
    TrainingArguments,
    Trainer,
)

processor = AutoImageProcessor.from_pretrained("shi-labs/dinat-mini-in1k-224")
model = DinatForImageClassification.from_pretrained(
    "shi-labs/dinat-mini-in1k-224",
    num_labels=10,
    ignore_mismatched_sizes=True,  # 替换分类头以匹配新类别数
)

# 这里假设你已经用 AutoImageProcessor 将数据集预处理为 input ids 形式
training_args = TrainingArguments(
    output_dir="./dinat-finetuned",
    per_device_train_batch_size=16,
    learning_rate=5e-5,
    num_train_epochs=3,
    remove_unused_columns=False,
)

trainer = Trainer(model=model, args=training_args, train_dataset=train_ds)
trainer.train()

动手前仍要牢记两个前提：环境需已安装 NATTEN；输入图应经过 AutoImageProcessor 的 resize/归一化处理。更完整的训练管线可参考 图像分类任务指南 与 image-classification 示例脚本目录。

源码与测试导览

若想深入核对上述行为，建议按以下路径阅读：

• 模型配置：DinatConfig 全部字段与自动推导逻辑；
• 模型实现：DinatModel（主干 + 池化）、DinatForImageClassification（分类头）、DinatBackbone（多尺度特征）、以及 NeighborhoodAttention 对 NATTEN 算子的调用；
• 模型测试：测试通过 DinatModelTester 以 embed_dim=16、kernel_size=3、dilations=[[3],[1,2],[1]] 等小配置跑通前向/反向与各类输出，并混合使用 BackboneTesterMixin、PipelineTesterMixin 验证 backbone 接口与 pipeline 兼容性；
• 依赖探测：is_natten_available() 与 NATTEN_IMPORT_ERROR 是缺失 NATTEN 时报错的来源。

总结

DiNAT 的价值在于把"局部邻域注意力 + 稀疏膨胀注意力"组织成分层 Transformer，在几乎不增加计算量的前提下补回了长程依赖建模能力，并在官方仓库中得到了完整的配置、主干、分类头与 backbone 四种落地形态。使用它只需把握三个要点：先装好 NATTEN、保持 patch size 为 4、区分 channels-last 的 hidden_states 与 channels-first 的 reshaped_hidden_states。需要多尺度金字塔特征的下游任务（如目标检测、语义分割）则应优先选用 DinatBackbone，配合 out_features 直接获取各 stage 的特征图。