十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

如何复现 Vision Transformer 论文的 patch embedding:用 nn.Conv2d 给图像分块并用 Flatten 展平?

如何复现 Vision Transformer 论文的 patch embedding:用 nn.Conv2d 给图像分块并用 Flatten 展平? 如何复现 Vision Transformer 论文的 patch embedding用 nn.Conv2d 给图像分块并用 Flatten 展平【免费下载链接】pytorch-deep-learningMaterials for the Learn PyTorch for Deep Learning: Zero to Mastery course.项目地址: https://gitcode.com/GitHub_Trending/py/pytorch-deep-learning在pytorch-deep-learningLearn PyTorch for Deep Learning 课程的论文复现项目中Milestone Project 2 要做的第一件事就是把 ViT 论文 Equation 1 里的 patch embedding 用 PyTorch 从零写出来。本文对应 08. PyTorch Paper Replicating 中第 4 节的 4.34.5 小节用nn.Conv2d()把一张 224×224 的三通道图像切成 16×16 的分块再用nn.Flatten()把分块特征图展平成[batch_size, N, P^2•C]的序列最后封装成一个可复用的PatchEmbedding层。完成后的可核对目标是输入(1, 3, 224, 224)输出形状为(1, 196, 768)。准备条件环境版本、数据与前置脚本notebook 第 0 节明确要求本 notebook 需要torch 1.12和torchvision 0.13notebook 内会先断言版本不满足时安装 nightly 版本在 Google Colab 中若触发了安装需要重启 runtime 后重跑该单元格# For this notebook to run with updated APIs, we need torch 1.12 and torchvision 0.13 try: import torch import torchvision assert int(torch.__version__.split(.)[1]) 12 or int(torch.__version__.split(.)[0]) 2, torch version should be 1.12 assert int(torchvision.__version__.split(.)[1]) 13, torchvision version should be 0.13 print(ftorch version: {torch.__version__}) print(ftorchvision version: {torchvision.__version__}) except: print(f[INFO] torch/torchvision versions not as required, installing nightly versions.) !pip3 install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 import torch import torchvision print(ftorch version: {torch.__version__}) print(ftorchvision version: {torchvision.__version__})数据沿用课程一直使用的 FoodVision Minipizza、steak、sushi 三类图像。仓库自带数据源文件 pizza_steak_sushi.zip也可以用 notebook 里的download_data()来自 helper_functions.py从远端下载并解压到pizza_steak_sushi/。之后按 ViT 论文 Table 3 的 Training resolution 224 构造 transforms 和DataLoaderimport torch import torchvision from torch import nn from torchvision import transforms from going_modular.going_modular import data_setup from helper_functions import download_data, set_seeds # Download pizza, steak, sushi images from GitHub image_path download_data(sourcehttps://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip, destinationpizza_steak_sushi) train_dir image_path / train test_dir image_path / test # Create image size (from Table 3 in the ViT paper) IMG_SIZE 224 manual_transforms transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), ]) train_dataloader, test_dataloader data_setup.create_dataloaders( train_dirtrain_dir, test_dirtest_dir, transformsmanual_transforms, batch_size32, ) # Get a single image and label from a batch image_batch, label_batch next(iter(train_dataloader)) image, label image_batch[0], label_batch[0]注意两个细节由于是从头训练不做 transfer learningtransforms 里没有提供normalizecreate_dataloaders()使用了pin_memoryTrue参数。本任务只需要单张image形状(3, 224, 224)不需要真的训练模型。先手算期望形状作为最后的验证基准在写层之前notebook 4.1 小节先按 ViT 论文 section 3.1 的记号手算输入输出形状后面所有代码都用它来判断对错。取 ViT-Base 的patch_size16、训练分辨率 224、三通道图像# Create example values height 224 # H (The training resolution is 224.) width 224 # W color_channels 3 # C patch_size 16 # P # Calculate N (number of patches) number_of_patches int((height * width) / patch_size**2) print(fNumber of patches (N) with image height (H{height}), width (W{width}) and patch size (P{patch_size}): {number_of_patches}) # Input shape (this is the size of a single image) embedding_layer_input_shape (height, width, color_channels) # Output shape embedding_layer_output_shape (number_of_patches, patch_size**2 * color_channels) print(fInput shape (single 2D image): {embedding_layer_input_shape}) print(fOutput shape (single 2D image flattened into patches): {embedding_layer_output_shape})文档示例输出N 196输入形状(224, 224, 3)输出形状(196, 768)。其中 196 H·W / P² 是 patch 数量也是后续 Transformer 的有效输入序列长度768 P²·C论文 Table 1 中 ViT-Base 的 Hidden size D 也是 768。这两个数字就是全文的验证基准patch embedding 层必须把(H, W, C)变成N × (P²·C)。用 nn.Conv2d 分块kernel_size 和 stride 都设为 patch_sizeViT 论文 section 3.1 说明 patch embedding 可以用卷积实现把卷积层的kernel_size和stride都设为patch_size卷积核就会以不重叠的方式扫过整张图像每经过一个 patch 就产出一个可学习的投影向量论文里叫 Linear Projection。out_channels设为嵌入维度 DViT-Base 为 768# Set the patch size patch_size 16 # Create the Conv2d layer with hyperparameters from the ViT paper conv2d nn.Conv2d(in_channels3, # number of color channels out_channels768, # from Table 1: Hidden size D, this is the embedding size kernel_sizepatch_size, # could also use (patch_size, patch_size) stridepatch_size, padding0)把单张图像送进去时nn.Conv2d要求输入带 batch 维所以用image.unsqueeze(0)# Pass the image through the convolutional layer image_out_of_conv conv2d(image.unsqueeze(0)) # add a single batch dimension (height, width, color_channels) - (batch, height, width, color_channels) print(image_out_of_conv.shape)文档示例输出为torch.Size([1, 768, 14, 14])读法是[batch_size, embedding_dim, feature_map_height, feature_map_width]224/16 14所以得到 14×14 个位置、每个位置一个 768 维特征。检查点如果打印出来的不是这个形状说明kernel_size/stride没有都设成patch_size或图像不是 224×224。用 nn.Flatten 展平特征图的空间维此时输出里已经含有 768 这个嵌入维度但 14×14 的空间维还是 2D 的而目标输出是 1 维的 196 个 patch 序列。notebook 4.4 小节用nn.Flatten()解决关键在于start_dim/end_dim只展平特征图的空间维维度 2 和 3不要动 batch 和 embedding 维# Create flatten layer flatten nn.Flatten(start_dim2, # flatten feature_map_height (dimension 2) end_dim3) # flatten feature_map_width (dimension 3) # 1. Take a single image, put it through the convolutional layer image_out_of_conv conv2d(image.unsqueeze(0)) # add batch dimension to avoid shape errors print(fImage feature map shape: {image_out_of_conv.shape}) # 2. Flatten the feature maps image_out_of_conv_flattened flatten(image_out_of_conv) print(fFlattened image feature map shape: {image_out_of_conv_flattened.shape})文档示例输出为(1, 768, 196)。它和目标形状(196, 768)只差维度顺序用permute(0, 2, 1)调整# Get flattened image patch embeddings in right shape image_out_of_conv_flattened_reshaped image_out_of_conv_flattened.permute(0, 2, 1) # [batch_size, P^2•C, N] - [batch_size, N, P^2•C] print(fPatch embedding sequence shape: {image_out_of_conv_flattened_reshaped.shape} - [batch_size, num_patches, embedding_size])文档示例输出为torch.Size([1, 196, 768])与手算的N × (P²·C)一致至此分块 展平这条链路已经跑通。封装成 PatchEmbedding 模块4.5 小节把上面的三步收进一个继承nn.Module的类参数默认值对应 ViT-Basein_channels3、patch_size16、embedding_dim768。forward()里还有一个断言保证输入图像边长能被patch_size整除# 1. Create a class which subclasses nn.Module class PatchEmbedding(nn.Module): Turns a 2D input image into a 1D sequence learnable embedding vector. Args: in_channels (int): Number of color channels for the input images. Defaults to 3. patch_size (int): Size of patches to convert input image into. Defaults to 16. embedding_dim (int): Size of embedding to turn image into. Defaults to 768. # 2. Initialize the class with appropriate variables def __init__(self, in_channels:int3, patch_size:int16, embedding_dim:int768): super().__init__() # 3. Create a layer to turn an image into patches self.patcher nn.Conv2d(in_channelsin_channels, out_channelsembedding_dim, kernel_sizepatch_size, stridepatch_size, padding0) # 4. Create a layer to flatten the patch feature maps into a single dimension self.flatten nn.Flatten(start_dim2, # only flatten the feature map dimensions into a single vector end_dim3) # 5. Define the forward method def forward(self, x): # Create assertion to check that inputs are the correct shape image_resolution x.shape[-1] assert image_resolution % patch_size 0, fInput image size must be divisible by patch size, image shape: {image_resolution}, patch size: {patch_size} # Perform the forward pass x_patched self.patcher(x) x_flattened self.flatten(x_patched) # 6. Make sure the output shape has the right order return x_flattened.permute(0, 2, 1) # adjust so the embedding is on the final dimension [batch_size, P^2•C, N] - [batch_size, N, P^2•C]对单张图像做端到端验证记得同样要先unsqueeze(0)加 batch 维否则nn.Conv2d会直接报错set_seeds() # Create an instance of patch embedding layer patchify PatchEmbedding(in_channels3, patch_size16, embedding_dim768) # Pass a single image through print(fInput image shape: {image.unsqueeze(0).shape}) patch_embedded_image patchify(image.unsqueeze(0)) # add an extra batch dimension on the 0th index, otherwise will error print(fOutput patch embedding shape: {patch_embedded_image.shape})文档示例输出输入torch.Size([1, 3, 224, 224])输出torch.Size([1, 196, 768])——与 4.1 小节手算的N196、P²•C768完全一致说明 patch embedding 部分复现完成。可选用 torchinfo 查看层的输入输出概要notebook 中安装torchinfo后可用summary快速核对层的输入/输出形状notebook 原代码是注释掉的取消注释即可运行把input_size换成下面注释里的错误尺寸可以看到断言失败的效果# Create random input sizes random_input_image (1, 3, 224, 224) random_input_image_error (1, 3, 250, 250) # will error because image size is incompatible with patch_size # Get a summary of the input and outputs of PatchEmbedding (uncomment for full output) # summary(PatchEmbedding(), # input_sizerandom_input_image, # try swapping this for random_input_image_error # col_names[input_size, output_size, num_params, trainable], # col_width20, # row_settings[var_names])限制与下一步输入必须是 4Dnn.Conv2d不接受(3, 224, 224)的单图必须unsqueeze(0)得到(1, 3, 224, 224)这是 notebook 中两处都特别注释过的报错点。图像边长必须能被 patch_size 整除forward()里的断言会在边长不满足时抛出Input image size must be divisible by patch size错误如 250×250 配patch_size16。这一层只覆盖 Equation 1 的一部分可学习的 class token 和 position embedding 不在这个模块里。文档中 class token 是用nn.Parameter(torch.ones(batch_size, 1, embedding_dimension))演示创建、再torch.cat(..., dim1)前置到 patch 序列开头torch.ones()仅为演示实际训练通常用torch.randn()初始化这部分继续看 08. PyTorch Paper Replicating 的 4.6 小节。形状口径本文所有(1, 768, 14, 14)、(1, 196, 768)等均为 notebook 的文档示例输出embedding_dim换用 Table 1 的其他 ViT 变体如 ViT-Large 的 1024时输出第二维会相应变化patch 数量 196 不变。【免费下载链接】pytorch-deep-learningMaterials for the Learn PyTorch for Deep Learning: Zero to Mastery course.项目地址: https://gitcode.com/GitHub_Trending/py/pytorch-deep-learning创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表