在DeiT里面，是通过mask的方式，将mask+unmasked的patches输出进ViT中，但其实在下游任务输入的patches还是和训练时patches的数量N是一致的（encoder所有的patches）。

而MAE是在encoder中只encoder未被mask的patches

通过什么方式支持的？

• 在处理文本时，可以根据最长的句子在批次中动态padding或截断长句子
• 而在处理图像（如使用ViT）时，可以将图像划分为大小相等的patches，数量可以根据图像的大小动态变化。

在训练阶段，部分patches被mask为0，但是处理的所有patches加起来的总长度还是一样的。被mask的位置在模型内部仍然占位，保持了输入序列的“框架”。这样，即使实际参与计算的只是部分元素，模型也能够适应在推理时使用全部元素的情况。

具体的计算步骤如下：

1. 确定mask哪些patches
2. 将mask的patches位置设置为0
3. 这些被mask和未被mask的所有patches一起被输入进attention模块
4. 将被mask的patches的注意力分数手动设置为“无穷大负数”（-inf）
5. 这些被mask的patches的softmax值就会变为0，也就意味着这些patches并未参与注意力的计算

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


class MaskedSelfAttention(nn.Module):
    def __init__(self, embed_size):
        super(MaskedSelfAttention, self).__init__()
        self.query = nn.Linear(embed_size, embed_size)
        self.key = nn.Linear(embed_size, embed_size)
        self.value = nn.Linear(embed_size, embed_size)

    def forward(self, x, mask=None):
        Q = self.query(x)
        K = self.key(x)
        V = self.value(x)

        # 计算自注意力得分
        attention_scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(Q.size(-1), dtype=torch.float32))

        # 将mask值为0的位置在attention_scores中设置为一个非常大的负数
        attention_scores = attention_scores.masked_fill(mask == 0, float('-inf'))

        # 使得这些位置的softmax结果接近0
        attention_weights = F.softmax(attention_scores, dim=-1)
        
        # 算最终的注意力加权和
        output = torch.matmul(attention_weights, V)

        return output


# 假设嵌入大小为512
embed_size = 512
# 创建一个mask，假设我们有4个patches，我们想要mask掉第2个和第4个patches
mask = torch.tensor([[1, 0, 1, 0]])
# 扩展mask维度以适应attention_scores的形状（假设批大小为1，序列长度为4），mask需要与attention_scores形状匹配，即(batch_size, 1, 1, seq_length)
mask = mask.unsqueeze(1).unsqueeze(2)

# 初始化模型和数据
sa = MaskedSelfAttention(embed_size)
x = torch.randn(1, 4, embed_size)  # 假设有一个批大小为1，序列长度为4的输入

output = sa(x, mask)
print(output)