Dilated Neighborhood Attention Transformer

Dilated Neighborhood Attention Transformer是一种基于Neighborhood Attention Transformer的改进模型，它通过引入空洞卷积（Dilated Convolution）来扩大感受野，从而提高模型的性能。具体来说，Dilated Neighborhood Attention Transformer在每个层级中使用了多个不同的空洞卷积核，这些卷积核的空洞率逐渐增加，从而使得每个query的感受野逐渐扩大。这种方法可以在不增加计算复杂度的情况下提高模型的性能，特别是在处理长序列数据时效果更为明显。

以下是Dilated Neighborhood Attention Transformer的实现代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class DilatedNeighborhoodAttention(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, dilation_rate):
        super(DilatedNeighborhoodAttention, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=dilation_rate*(kernel_size-1), dilation=dilation_rate)
        self.norm = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        x = self.conv(x)
        x = self.norm(x)
        x = self.relu(x)
        return x

class DilatedNeighborhoodAttentionTransformer(nn.Module):
    def __init__(self, num_layers, num_heads, d_model, d_ff, dropout):
        super(DilatedNeighborhoodAttentionTransformer, self).__init__()
        self.num_layers = num_layers
        self.self_attentions = nn.ModuleList([nn.MultiheadAttention(d_model, num_heads, dropout=dropout) for _ in range(num_layers)])
        self.dilated_attentions = nn.ModuleList([DilatedNeighborhoodAttention(d_model, d_model, kernel_size=3, dilation_rate=2**i) for i in range(num_layers)])
        self.ffns = nn.ModuleList([nn.Sequential(nn.Linear(d_model, d_ff), nn.ReLU(inplace=True), nn.Linear(d_ff, d_model)) for _ in range(num_layers)])
        self.norms1 = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(num_layers)])
        self.norms2 = nn.ModuleList([nn.LayerNorm(d_model) for _ in range(num_layers)])
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        for i in range(self.num_layers):
            residual = x
            x, _ = self.self_attentions[i](x, x, x)
            x = self.norms1[i](residual + self.dropout(x))
            residual = x
            x = self.dilated_attentions[i](x)
            x = self.norms2[i](residual + self.dropout(x))
            residual = x
            x = self.ffns[i](x)
            x = self.norms2[i](residual + self.dropout(x))
        return x