MAE的mask的代码逻辑

MAE（Masked Autoencoder for Image Recognition），是一种基于自编码器架构的无监督预训练模型，它通过掩码掉输入图像的一部分来学习图像的表示。在PyTorch中实现MAE的mask逻辑通常涉及以下几个步骤：

1. **创建输入样本**: 将原始图片随机选择部分像素设置为0，即进行遮罩操作。这通常通过`torch.nn.functional.embedding`函数完成，其中嵌入矩阵对应于全0和全1两种情况。

import torch
from einops import rearrange

def create_mask(input_img, mask_prob=0.75):
    binary_mask = torch.bernoulli(torch.full_like(input_img, p=mask_prob))
    return rearrange(binary_mask, 'b c h w -> b (h w) c')

2. **应用mask**: 创建的掩码会被应用于原图，生成带掩码的输入图像 (`masked_input`) 和对应的掩码 (`mask`), 这些都会用于训练过程中恢复原始图像的重建任务。

input_img = ... # 输入图片
mask = create_mask(input_img)
masked_input = input_img * mask + (1 - mask) * 0  # 或者用mean等填充值替换0

3. **模型前向传播**: 遮罩后的图片传递给MAE模型进行训练，模型会尝试从这个部分信息丢失的输入中重构出原始图像。

4. **损失计算**: 通常使用均方误差（MSE）或其他适合图像恢复的任务的损失函数来评估模型的表现。

reconstructed_img = model(masked_input)
loss = F.mse_loss(reconstructed_img, input_img, reduction='none') * mask