def Bottlenecktransformer

BoTNet（Bottleneck Transformer Network）是一种基于Transformer的新骨干架构，它同时使用卷积和自注意力机制。在ResNet的最后3个bottleneck blocks中使用全局多头自注意力（Multi-Head Self-Attention, MHSA）替换3 × 3空间卷积。这种结构可以在保持高精度的同时，大大减少模型的计算量和参数数量。BoTNet已经在多个计算机视觉任务中取得了优异的表现，例如图像分类、目标检测和语义分割等。

import torch
import torch.nn as nn

class MHSA(nn.Module):
    def __init__(self, dim, num_heads):
        super(MHSA, self).__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim ** -0.5
        self.query = nn.Linear(dim, dim, bias=False)
        self.key = nn.Linear(dim, dim, bias=False)
        self.value = nn.Linear(dim, dim, bias=False)
        self.out = nn.Linear(dim, dim)

    def forward(self, x):
        b, n, _, h = *x.shape, self.num_heads
        qkv = [l(x).reshape(b, n, h, self.head_dim).transpose(1, 2) for l, x in zip((self.query, self.key, self.value), (x, x, x))]
        dots = qkv[0] @ qkv[1].transpose(-2, -1) * self.scale
        attn = dots.softmax(dim=-1)
        out = attn @ qkv[2]
        out = out.transpose(1, 2).reshape(b, n, -1)
        return self.out(out)

class Bottleneck(nn.Module):
    def __init__(self, dim, num_heads, expansion_factor=4):
        super(Bottleneck, self).__init__()
        self.mhsa = MHSA(dim, num_heads)
        self.norm1 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, expansion_factor * dim),
            nn.GELU(),
            nn.Linear(expansion_factor * dim, dim),
        )
        self.norm2 = nn.LayerNorm(dim)

    def forward(self, x):
        x = x + self.mhsa(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class BoTNet(nn.Module):
    def __init__(self, layers, channels, num_classes, num_heads=4, expansion_factor=4):
        super(BoTNet, self).__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, channels[0], kernel_size=3, stride=2, padding=1, bias=False),
            nn.BatchNorm2d(channels[0]),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(channels[0]),
            nn.ReLU(inplace=True),
            nn.Conv2d(channels[0], channels[0], kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(channels[0]),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
        )
        self.layer1 = self._make_layer(dim=channels[0], num_heads=num_heads, expansion_factor=expansion_factor, num_blocks=layers[0])
        self.layer2 = self._make_layer(dim=channels[1], num_heads=num_heads, expansion_factor=expansion_factor, num_blocks=layers[1], stride=2)
        self.layer3 = self._make_layer(dim=channels[2], num_heads=num_heads, expansion_factor=expansion_factor, num_blocks=layers[2], stride=2)
        self.layer4 = self._make_layer(dim=channels[3], num_heads=num_heads, expansion_factor=expansion_factor, num_blocks=layers[3], stride=2)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(channels[3], num_classes)

    def _make_layer(self, dim, num_heads, expansion_factor, num_blocks, stride=1):
        layers = [Bottleneck(dim, num_heads, expansion_factor=expansion_factor)]
        for _ in range(num_blocks - 1):
            layers.append(Bottleneck(dim, num_heads, expansion_factor=expansion_factor))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.stem(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.flatten(1)
        x = self.fc(x)
        return x