class Bottleneck(nn.Module):
    # 定义扩展倍数，用于决定输出通道数的扩展比例
    expansion = 4

    def __init__(self, inplanes, planes, stride=1):
        """
        Bottleneck 模块的初始化方法。
        参数：
        - inplanes: 输入通道数。
        - planes: 每一层卷积的基础通道数，最终会扩展为 planes * expansion。
        - stride: 步幅，控制空间分辨率的缩小（默认值为1）。
        """
        super().__init__()

        # 第1个卷积层，使用 1x1 卷积，用于减少计算量
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)  # 批归一化
        self.relu1 = nn.ReLU(inplace=True)  # 激活函数，inplace=True 节省内存

        # 第2个卷积层，使用 3x3 卷积，用于提取特征
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)  # 批归一化
        self.relu2 = nn.ReLU(inplace=True)  # 激活函数

        # 如果 stride > 1，使用 AvgPool2d 进行空间降采样，否则使用 Identity（保持不变）
        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()

        # 第3个卷积层，使用 1x1 卷积，用于扩展通道数至 planes * expansion
        self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * self.expansion)  # 批归一化
        self.relu3 = nn.ReLU(inplace=True)  # 激活函数

        # 如果输入和输出的通道数或分辨率不同，则需要通过 downsample 进行调整
        self.downsample = None
        self.stride = stride

        if stride > 1 or inplanes != planes * Bottleneck.expansion:
            # downsample 模块，用于调整 shortcut 的形状，使其与主路径的输出形状匹配
            self.downsample = nn.Sequential(OrderedDict([
                # 使用 AvgPool2d 进行空间降采样（仅在 stride > 1 时）
                ("-1", nn.AvgPool2d(stride)),
                # 1x1 卷积调整通道数
                ("0", nn.Conv2d(inplanes, planes * self.expansion, kernel_size=1, stride=1, bias=False)),
                # 批归一化
                ("1", nn.BatchNorm2d(planes * self.expansion))
            ]))

    def forward(self, x: torch.Tensor):
        """
        前向传播函数。
        参数：
        - x: 输入张量，形状为 (batch_size, inplanes, H, W)
        返回值：
        - 输出张量，形状为 (batch_size, planes * expansion, H / stride, W / stride)
        """
        identity = x  # 保存输入作为 shortcut（残差连接的输入）

        # 主路径：第1个卷积 + BN + ReLU
        out = self.relu1(self.bn1(self.conv1(x)))
        # 主路径：第2个卷积 + BN + ReLU
        out = self.relu2(self.bn2(self.conv2(out)))
        # 如果 stride > 1，进行空间降采样
        out = self.avgpool(out)
        # 主路径：第3个卷积 + BN
        out = self.bn3(self.conv3(out))

        # 如果 downsample 不为 None，则对 identity 进行调整
        if self.downsample is not None:
            identity = self.downsample(x)

        # 残差连接：将主路径的输出和 shortcut 相加
        out += identity
        # 激活函数
        out = self.relu3(out)
        return out

class AttentionPool2d(nn.Module):
    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
        """
        初始化 Attention Pooling 模块。

        参数：
        - spacial_dim (int): 输入特征图的空间维度（即特征图宽高的大小，假设特征图是正方形）。
        - embed_dim (int): 输入特征的嵌入维度（通道数）。
        - num_heads (int): 多头注意力的头数。
        - output_dim (int, optional): 输出特征的目标维度。如果未指定，则默认为 embed_dim。
        """
        super().__init__()
        # 位置嵌入 (Positional Embedding)，初始化为随机值，大小为 (spacial_dim^2 + 1, embed_dim)
        # 额外的 "+1" 是为全局 token 预留的位置。
        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)

        # Query、Key、Value 的线性投影
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)

        # 输出投影，将嵌入维度从 embed_dim 调整为 output_dim
        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)

        # 多头注意力的头数
        self.num_heads = num_heads

    def forward(self, x):
        """
        前向传播函数，执行 Attention Pooling 操作。

        参数：
        - x (Tensor): 输入特征图，形状为 (batch_size, channels, height, width)。

        返回：
        - Tensor: 输出的全局特征向量，形状为 (batch_size, output_dim)。
        """
        # 1. 展平特征图并转置
        # 原始形状：x -> (batch_size, channels, height, width)
        # 转换后：x -> (height * width, batch_size, channels)
        x = x.flatten(start_dim=2).permute(2, 0, 1)

        # 2. 添加全局 token
        # 全局 token 是整幅特征图的均值，用于捕获全局信息。
        # `x.mean(dim=0, keepdim=True)` 的形状为 (1, batch_size, channels)
        # 添加到特征图的前面，新的 x 形状为 (height * width + 1, batch_size, channels)
        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)

        # 3. 添加位置嵌入
        # self.positional_embedding 的形状为 (height * width + 1, embed_dim)
        # 将位置嵌入添加到特征图上，每个位置都会加上对应的嵌入值。
        x = x + self.positional_embedding[:, None, :].to(x.dtype)

        # 4. 执行多头注意力
        # 注意力操作的参数：
        # - query=x[:1]: 使用全局 token 作为 Query，形状为 (1, batch_size, embed_dim)。
        # - key=x: 使用整个特征图（包括全局 token）作为 Key，形状为 (height * width + 1, batch_size, embed_dim)。
        # - value=x: 同上，使用整个特征图作为 Value。
        # 注意力机制的输出形状为 (1, batch_size, embed_dim)，即仅保留全局 token 的注意力输出。
        x, _ = F.multi_head_attention_forward(
            query=x[:1],  # 全局 token 作为 Query
            key=x,  # 整个特征图作为 Key
            value=x,  # 整个特征图作为 Value
            embed_dim_to_check=x.shape[-1],  # 嵌入维度
            num_heads=self.num_heads,  # 注意力头的数量
            q_proj_weight=self.q_proj.weight,  # Query 的投影权重
            k_proj_weight=self.k_proj.weight,  # Key 的投影权重
            v_proj_weight=self.v_proj.weight,  # Value 的投影权重
            in_proj_weight=None,  # 不使用合并的投影权重
            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),  # Query, Key, Value 的偏置
            bias_k=None,  # Key 的额外偏置（未使用）
            bias_v=None,  # Value 的额外偏置（未使用）
            add_zero_attn=False,  # 不添加额外的全零注意力头
            dropout_p=0,  # 不使用 Dropout
            out_proj_weight=self.c_proj.weight,  # 输出投影的权重
            out_proj_bias=self.c_proj.bias,  # 输出投影的偏置
            use_separate_proj_weight=True,  # 使用单独的 Query, Key, Value 投影权重
            training=self.training,  # 根据当前模式确定是否训练
            need_weights=False  # 不需要返回注意力权重
        )

        # 5. 移除第一个维度
        # 注意力输出的形状为 (1, batch_size, embed_dim)，去掉第一个维度后，变为 (batch_size, embed_dim)。
        return x.squeeze(0)

class ModifiedResNet(nn.Module):
    """
    一个基于 ResNet 的改进版本，包含以下变化：
    - 起始部分使用了 3 层卷积（stem），而不是原始 ResNet 的单层卷积，同时将 max pool 替换为 avg pool。
    - 对于 stride > 1 的卷积层，使用了反锯齿（anti-aliasing）的卷积方式，即在卷积前添加了 avg pool。
    - 最终的全局池化层替换为 QKV 注意力池化层，而不是传统的 average pool。
    """

    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
        """
        初始化函数。
        参数：
        - layers: 每个阶段的残差块数量（如 [3, 4, 6, 3] 表示每个阶段的残差块数）。
        - output_dim: 输出特征的维度。
        - heads: 注意力池化的头数。
        - input_resolution: 输入图像的分辨率（默认 224）。
        - width: 初始通道宽度（默认 64）。
        """
        super().__init__()
        self.output_dim = output_dim
        self.input_resolution = input_resolution

        # 初始化阶段（stem）：3 层卷积代替原来的单层卷积
        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)  # 输入通道为3，输出为 width//2
        self.bn1 = nn.BatchNorm2d(width // 2)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(width // 2)
        self.relu2 = nn.ReLU(inplace=True)
        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)  # 输出为 width
        self.bn3 = nn.BatchNorm2d(width)
        self.relu3 = nn.ReLU(inplace=True)
        self.avgpool = nn.AvgPool2d(2)  # 平均池化，用于下采样

        # 残差层
        self._inplanes = width  # 初始输入通道数，用于动态调整残差块的输入
        self.layer1 = self._make_layer(width, layers[0])  # 第一阶段残差块
        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)  # 第二阶段残差块，输出通道翻倍
        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)  # 第三阶段残差块
        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)  # 第四阶段残差块

        # 注意力池化层
        embed_dim = width * 32  # ResNet 的特征维度，等于最后一个阶段的输出通道数
        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)  # QKV 注意力池化

    def _make_layer(self, planes, blocks, stride=1):
        """
        创建一个残差模块（由多个 Bottleneck 组成）。
        参数：
        - planes: 当前阶段的通道数。
        - blocks: 残差块数量。
        - stride: 当前阶段的步幅。
        返回值：
        - 一个由多个残差块组成的序列模块。
        """
        layers = [Bottleneck(self._inplanes, planes, stride)]  # 第一个残差块，可能需要调整步幅
        self._inplanes = planes * Bottleneck.expansion  # 更新通道数
        for _ in range(1, blocks):
            layers.append(Bottleneck(self._inplanes, planes))  # 后续残差块保持步幅为 1

        return nn.Sequential(*layers)  # 返回残差块的序列

    def forward(self, x):
        """
        前向传播函数。
        参数：
        - x: 输入张量，形状为 (batch_size, channels, height, width)。
        返回值：
        - x: 输出张量，经过特征提取和注意力池化后。
        """
        def stem(x):
            """
            起始阶段（stem）：连续 3 层卷积 + 平均池化。
            参数：
            - x: 输入张量。
            返回值：
            - x: 经过 stem 阶段后的特征张量。
            """
            x = self.relu1(self.bn1(self.conv1(x)))  # 第1层卷积
            x = self.relu2(self.bn2(self.conv2(x)))  # 第2层卷积
            x = self.relu3(self.bn3(self.conv3(x)))  # 第3层卷积
            x = self.avgpool(x)  # 平均池化
            return x

        x = x.type(self.conv1.weight.dtype)  # 将输入类型调整为卷积权重的类型（支持混合精度训练）
        x = stem(x)  # 经过起始阶段
        x = self.layer1(x)  # 第一阶段残差层
        x = self.layer2(x)  # 第二阶段残差层
        x = self.layer3(x)  # 第三阶段残差层
        x = self.layer4(x)  # 第四阶段残差层
        x = self.attnpool(x)  # 最后通过注意力池化
        return x


# 实例化模型
model = ModifiedResNet(layers=[3, 4, 6, 3], output_dim=512, heads=8)

# 输入示例数据
input_data = torch.randn(10, 3, 224, 224)  # (batch_size, channels, height, width)
output = model(input_data)

print("输出形状:", output.shape)

import torch
import torch.nn as nn

class PatchEmbedding(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768):
        """
        初始化 PatchEmbedding 类，用于将输入图像划分为若干 Patch，并投影到固定的嵌入维度。
        
        参数：
        - img_size: 输入图像的分辨率（默认为 224）。
        - patch_size: 每个 Patch 的尺寸（默认为 16x16）。
        - in_channels: 输入图像的通道数（默认为 3，即 RGB 图像）。
        - embed_dim: 嵌入维度（Transformer 的输入特征维度，默认为 768）。
        """
        super().__init__()
        self.img_size = img_size  # 保存输入图像的尺寸
        self.patch_size = patch_size  # 保存每个 Patch 的尺寸
        self.num_patches = (img_size // patch_size) * (img_size // patch_size)  
        # 计算图像被划分的总 Patch 数量 = (图像宽度上的 Patch 数) * (图像高度上的 Patch 数)
        self.patch_dim = patch_size * patch_size * in_channels  
        # 每个 Patch 的维度 = patch_size * patch_size * 通道数（即展平后的特征数）

        # 定义一个线性层，将每个 Patch 的原始特征（patch_dim）投影到 Transformer 的嵌入维度（embed_dim）
        self.projection = nn.Linear(self.patch_dim, embed_dim)
        # 初始化一个可学习的参数，用于位置编码，形状为 (1, num_patches, embed_dim)
        self.position_embedding = nn.Parameter(torch.randn(1, self.num_patches, embed_dim))

    def forward(self, x):
        """
        前向传播函数，将输入图像转换为 Patch 嵌入，并添加位置编码。

        参数：
        - x: 输入张量，形状为 (batch_size, in_channels, height, width)。

        返回值：
        - embeddings: 包含位置编码的 Patch 嵌入，形状为 (batch_size, num_patches, embed_dim)。
        """
        # 获取输入张量的批量大小、通道数、高度和宽度
        B, C, H, W = x.shape
        # 确保输入图像的分辨率与定义的 img_size 一致
        assert H == self.img_size and W == self.img_size, "Input size must match img_size"
        
        # 将输入图像划分为 Patch，使用 unfold 函数：
        # unfold(2, patch_size, patch_size) 在高度维度上以 patch_size 的窗口滑动并提取小块；
        # unfold(3, patch_size, patch_size) 在宽度维度上以 patch_size 的窗口滑动并提取小块。
        patches = x.unfold(2, self.patch_size, self.patch_size).unfold(3, self.patch_size, self.patch_size)
        
        # contiguous() 确保数据在内存中是连续的；
        # view 将张量重塑为 (B, num_patches, patch_dim)，
        # 其中每个 Patch 被展平为长度为 patch_dim 的向量。
        patches = patches.contiguous().view(B, -1, self.patch_dim)

        # 使用线性层将每个 Patch 的特征投影到嵌入维度（embed_dim），
        # 输出的形状为 (B, num_patches, embed_dim)。
        embeddings = self.projection(patches)

        # 添加位置编码，位置编码的形状为 (1, num_patches, embed_dim)，
        # 会广播到 batch 的每个样本。
        embeddings += self.position_embedding

        # 返回添加了位置编码的 Patch 嵌入
        return embeddings

class LayerNorm(nn.LayerNorm):
    """
    LayerNorm 的子类，用于处理 fp16（半精度浮点数）。
    """

    def forward(self, x: torch.Tensor):
        # 保存原始数据类型
        orig_type = x.dtype
        # 将输入转换为 float32 类型，执行 LayerNorm 操作
        ret = super().forward(x.type(torch.float32))
        # 恢复原始数据类型
        return ret.type(orig_type)


class QuickGELU(nn.Module):
    """
    实现 QuickGELU 激活函数，比标准 GELU 更简单且计算效率更高。
    """

    def forward(self, x: torch.Tensor):
        # QuickGELU 激活函数公式：x * sigmoid(1.702 * x)
        return x * torch.sigmoid(1.702 * x)


class ResidualAttentionBlock(nn.Module):
    """
    实现一个残差注意力模块，包括多头注意力机制和前馈网络（MLP）。
    """

    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
        """
        初始化函数。
        参数：
        - d_model: 模型的隐藏层维度（即特征向量的维度）。
        - n_head: 多头注意力的头数。
        - attn_mask: 注意力掩码，可选。
        """
        super().__init__()

        # 多头注意力模块
        self.attn = nn.MultiheadAttention(d_model, n_head)
        # 第一个 LayerNorm
        self.ln_1 = LayerNorm(d_model)
        # 前馈网络（MLP），包括线性层、QuickGELU 激活和另一层线性变换
        self.mlp = nn.Sequential(OrderedDict([
            ("c_fc", nn.Linear(d_model, d_model * 4)),  # 扩展特征维度 4 倍
            ("gelu", QuickGELU()),  # QuickGELU 激活
            ("c_proj", nn.Linear(d_model * 4, d_model))  # 恢复到原始特征维度
        ]))
        # 第二![](https://skojiangdoc.oss-cn-beijing.aliyuncs.com/2024LLM/training/01.png)个 LayerNorm
        self.ln_2 = LayerNorm(d_model)
        # 注意力掩码
        self.attn_mask = attn_mask

    def attention(self, x: torch.Tensor):
        """
        计算多头注意力。
        """
        # 如果存在注意力掩码，将其转换为与输入相同的数据类型和设备
        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
        # 调用多头注意力，返回注意力结果（不需要权重）
        return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0]

    def forward(self, x: torch.Tensor):
        """
        前向传播。
        """
        # 1. 多头注意力：先进行 LayerNorm，再计算注意力，并与输入 x 相加（残差连接）
        x = x + self.attention(self.ln_1(x))
        # 2. 前馈网络：先进行 LayerNorm，再通过 MLP，并与输入 x 相加（残差连接）
        x = x + self.mlp(self.ln_2(x))
        return x


class Transformer(nn.Module):
    """
    基于 ResidualAttentionBlock 的 Transformer 模块。
    """

    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
        """
        初始化函数。
        参数：
        - width: 模型的隐藏层维度。
        - layers: Transformer 层数。
        - heads: 多头注意力的头数。
        - attn_mask: 注意力掩码。
        """
        super().__init__()
        self.width = width
        self.layers = layers
        # 堆叠多个 ResidualAttentionBlock，构成完整的 Transformer
        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])

    def forward(self, x: torch.Tensor):
        """
        前向传播。
        """
        return self.resblocks(x)


class VisionTransformer(nn.Module):
    """
    实现 Vision Transformer，用于处理图像输入。
    """

    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
        """
        初始化函数。
        参数：
        - input_resolution: 输入图像的分辨率（如 224）。
        - patch_size: 图像切片的尺寸（如 16x16）。
        - width: Transformer 的隐藏层维度。
        - layers: Transformer 层数。
        - heads: 多头注意力的头数。
        - output_dim: 输出特征维度。
        """
        super().__init__()
        self.input_resolution = input_resolution
        self.output_dim = output_dim
        # 图像切片卷积，将输入图像分割为若干 patch，并映射到 width 维度
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)

        # 类别嵌入，用于标识全局图像特征
        scale = width ** -0.5
        self.class_embedding = nn.Parameter(scale * torch.randn(width))
        # 位置嵌入，用于编码每个 patch 的位置
        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
        # 输入的 LayerNorm
        self.ln_pre = LayerNorm(width)

        # Transformer 编码器
        self.transformer = Transformer(width, layers, heads)

        # 输出的 LayerNorm 和投影层
        self.ln_post = LayerNorm(width)
        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))

    def forward(self, x: torch.Tensor):
        """
        前向传播。
        参数：
        - x: 输入图像张量，形状为 (batch_size, 3, H, W)。
        返回值：
        - x: 输出特征，形状为 (batch_size, output_dim)。
        """
        # 图像切片，将输入分割为 patch 并映射到特征维度
        x = self.conv1(x)  # 形状 = [batch_size, width, grid, grid]
        # 将每个 patch 展平，并调整为 [batch_size, num_patches, width]
        x = x.reshape(x.shape[0], x.shape[1], -1)  # 形状 = [batch_size, width, grid ** 2]
        x = x.permute(0, 2, 1)  # 调整为 [batch_size, num_patches, width]

        # 添加类别嵌入
        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # 形状 = [batch_size, num_patches + 1, width]
        # 加上位置嵌入
        x = x + self.positional_embedding.to(x.dtype)
        # 输入的 LayerNorm
        x = self.ln_pre(x)

        # Transformer 编码器
        x = x.permute(1, 0, 2)  # 调整为 [num_patches + 1, batch_size, width]
        x = self.transformer(x)
        x = x.permute(1, 0, 2)  # 调整回 [batch_size, num_patches + 1, width]

        # 输出特征仅取第一个（全局图像特征）
        x = self.ln_post(x[:, 0, :])

        # 如果存在投影层，将其映射到输出维度
        if self.proj is not None:
            x = x @ self.proj

        return x

# 调整参数以减少内存占用
input_resolution = 224  # 分辨率
width = 512  # 隐藏层维度
layers = 6  #  Transformer 层数
heads = 8  # 注意力头数
output_dim = 512  # 输出特征维度
batch_size = 10  # batch大小

# 重新定义 VisionTransformer 模型
vit = VisionTransformer(
    input_resolution=input_resolution,
    patch_size=16,
    width=width,
    layers=layers,
    heads=heads,
    output_dim=output_dim
)

# 假设输入图像数据 (10,3,224,224)
input_data = torch.randn(batch_size, 3, input_resolution, input_resolution)

# 通过模型前向传播
output = vit(input_data)

# 打印输出的形状
output.shape