当前位置：首页 > article >正文

对ViT 中Patch Embedding理解

article 2024/9/21 0:07:22

借鉴了这个博主的ViT Patch Embedding理解-CSDN博客，再加了一些理解。

就通过代码来理解吧

假设输入图像的维度为HxWxC，分别表示高，宽和通道数。

PatchEmbed 的类，它继承了 nn.Module，实现了将输入的2维图像（3通道）分割为多个小块（patches）(若干个不重叠的 patch)，并将每个小块映射到特定维度的嵌入（embedding）向量空间中。该类的核心思想是将输入的图像划分为固定大小的 patch，并通过卷积操作将这些 patch 转换为1维向量（embedding，经过编码的图像块）（线性变换）。

class PatchEmbed(nn.Module):
  """
    Image to Patch Embedding 得到的是经过编码的图片块
  """
 
  def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
    super().__init__()
    img_size = (img_size, img_size)
    patch_size = (patch_size, patch_size)
    num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
    self.img_size = img_size
    self.patch_size = patch_size
    self.num_patches = num_patches
 
    # 
    # embed_dim表示切好的图片拉成一维向量后的特征长度
    # 
    # 图像共切分为N = HW/P^2个patch块
    # 在实现上等同于对reshape后的patch序列进行一个PxP且stride为P的卷积操作
    # output = {[(n+2p-f)/s + 1]向下取整}^2
    # 即output = {[(n-P)/P + 1]向下取整}^2 = (n/P)^2
    # 
    self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
 
  def forward(self, x):
    B, C, H, W = x.shape
    assert H == self.img_size[0] and W == self.img_size[1], \
      f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
    x = self.proj(x).flatten(2).transpose(1, 2)
    return x  # x.shape is [8, 196, 768]

`1.init` 构造函数：

img_size=224: 输入图像的尺寸是 224x224。
patch_size=16: 将图像划分成 16x16 的小块（patch）。每个 patch 是一个三通道的小图像块。
in_chans=3: 表示输入图像的通道数为 3（通常为 RGB 图像）。
embed_dim=768: 每个 patch 最终会被映射到 768 维的向量空间（768组 3通道的卷积核，原来的一个图像块patch（3通道）映射成特征图的一个像素点（值），768 就是提取了768种特征然后将一个像素点映射到768维的特征空间）。输入的图像有 3 个通道（RGB 图像），每个卷积核都有 3 个通道的权重（滤波器）。每个 16x16 大小的 patch 经过一个卷积核，会被映射为一个单一的数值。每个 patch 被 768 个卷积核处理后得到了 768 个数值，这 768 个数值表示该 patch 在不同卷积核下的特征响应。这些特征数值构成了一个 768 维的向量，这意味着这个 patch 被映射到了 768 维的特征空间中。这些特征表示包括局部区域的颜色、纹理、边缘等信息，卷积核通过不断学习会提取出对任务有用的特征。
num_patches: 计算图像中有多少个 patch，即 (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1])。对于 224x224 的图像和 16x16 的 patch，它将产生 14x14 = 196 个 patch。

卷积层 (self.proj)

这一步使用一个二维卷积层来对图像进行 patch 的切分和 embedding 的生成。

kernel_size=patch_size: 卷积核的大小与 patch 的大小相同（16x16）。
stride=patch_size: 卷积的步长也是 16，因此卷积会以 16x16 的步幅滑动，即每次滑动的距离正好等于一个 patch 的大小。

这相当于将图像按块切分，并将每个 patch 通过卷积操作投影到 embed_dim 维的特征空间。

`2. forward` 函数

输入 x 的维度为 (B, C, H, W)，其中 B 是 batch size，C 是通道数（例如 3 个 RGB 通道），H 和 W 是图像的高度和宽度。
assert 语句确保输入的图像大小符合预期的尺寸 self.img_size，否则抛出异常。
self.proj(x) 通过卷积层将图像切分为 patch 并生成嵌入。
flatten(2) 将特征图的第三维和第四维（~~height 和 width~~，其实是那个(img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1])=num_patches）展平为一维，以便于后续处理。
transpose(1, 2) 交换维度，使得输出的形状为 [batch_size, num_patches, embed_dim（通道数线性变换了3→768）]，即每个 batch 中的每个 patch 都有一个 embed_dim 维的嵌入向量。

输出形状

输出的形状为 [8, 196, 768]：

8 表示 batch size。
196 表示图像被分割成 196 个 patch。
768 是每个 patch 被映射到的嵌入维度。

总结来说，这个类的功能是将输入图像通过卷积的方式分割成多个固定大小的 patch，并将每个 patch 转换为一个高维特征表示(其中 卷积核的主要作用是将图像切分成固定大小的 patch，同时也会进行一定的特征提取。这种操作不仅分割图像，还通过卷积层（768）对每个 patch 进行线性变换，将其映射到一个特征空间),用于后续处理。

卷积核通过学习不同的权重，能够提取出局部区域内的边缘、纹理、颜色等特征。

卷积核的初始化很重要，它会影响模型的收敛速度和最终效果。有几种常见的初始化方法.

Xavier 初始化（Glorot 初始化）：