掩码图像建模简介

L_1(\hat{y}, y) = \Vert \hat{y} - y \Vert_1

L_2(\hat{y}, y) = \Vert \hat{y} - y \Vert_2^2

\mathcal{E}_i = [e_i^1, e_i^2, ..., e_i^P] = \text{PatchEmbed}([ X_i^1, X_i^2, ..., X_i^P] )

\mathcal{E}_i \leftarrow [e^{\text{[CLS]}}, e_i^1, e_i^2, ..., e_i^P]

\mathcal{E}_i \leftarrow \mathcal{E}_i + \mathcal{S}

\mathcal{E}^q_i = \mathcal{E}_i W_q, \ \mathcal{E}^k_i = \mathcal{E}_i W_k, \ \mathcal{E}^v_i = \mathcal{E}_i W_v

\mathcal{A}_i = \text{softmax} \left ( \frac{\mathcal{E}^q_i (\mathcal{E}^k_i)^T}{\sqrt{l}} \right )

\mathcal{E}^o_i = \mathcal{A}_i \mathcal{E}^v_i

1import math
2import torch
3import torch.nn as nn
4from timm.models.vision_transformer import Block
5from PIL import Image
6import torchvision.transforms.functional as T
7import torchvision
8from utils import extract_image_patches
9
10class PatchEmbed(nn.Module):
11    """ 将图像转换为图像块嵌入
12    """
13    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
14        super().__init__()
15        num_patches = (img_size // patch_size) * (img_size // patch_size)
16        self.img_size = img_size
17        self.patch_size = patch_size
18        self.num_patches = num_patches
19
20        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
21            
22    def forward(self, x):
23        x = self.proj(x)
24        return x.flatten(2).transpose(1, 2)
25
26class VisionTransformer(nn.Module):
27    """ 视觉Transformer模型 """
28    def __init__(self, img_size=[224], patch_size=16, 
29                       in_chans=3, num_classes=0, 
30                       embed_dim=768, depth=12, num_heads=12):
31        super().__init__()
32        self.num_features = self.embed_dim = embed_dim
33        print("输入图像的形状:", img_size[0])
34        self.patch_embed = PatchEmbed(
35            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
36        
37        num_patches = self.patch_embed.num_patches
38        print("图像块总数 :", num_patches)
39
40        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
41        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
42        self.pos_drop = nn.Dropout(p=0)
43        print("位置编码的形状: ", self.pos_embed.shape)
44
45        self.blocks = nn.ModuleList([
46            Block(dim=embed_dim, num_heads=num_heads)
47            for i in range(depth)])
48
49        self.norm = nn.LayerNorm(embed_dim)
50        
51    def forward(self, x):
52        B, nc, w, h = x.shape
53        # 图像块线性嵌入
54        x = self.patch_embed(x)
55        print("图像块嵌入的形状:", x.shape)
56
57        # 将[CLS]标记添加到图像块嵌入中
58        print("【CLS】标记的形状 :", self.cls_token.shape)
59        cls_tokens = self.cls_token.expand(B, -1, -1)
60        x = torch.cat((cls_tokens, x), dim=1)
61        print("添加【CLS】标记后的形状 :", x.shape)
62
63        # 为每个标记添加位置编码
64        x = x + self.pos_embed
65
66        x = self.pos_drop(x)
67      
68        for blk in self.blocks: # 逐层通过多头自注意力模块
69            x = blk(x)
70        print("经过MSA层后的输出形状 :", x.shape)
71
72        x = self.norm(x) # 应用层归一化
73        return x[:, 0] # 返回【CLS】标记的表示
74
75model = VisionTransformer(patch_size=16, embed_dim=192, depth=12, num_heads=3)
76image  = T.to_tensor(Image.open("cat.jpg").resize((224,224)))[None, :, :, :]
77torchvision.utils.save_image(image, "./output/image.png", normalize=True)
78patches = extract_image_patches(image, kernel=16, stride=16)
79print(patches.shape)
80torchvision.utils.save_image(patches[0], "./output/patches.png", normalize=True, nrow=14)
81f = model(image)
82print("【CLS】标记特征的形状: ", f.shape)