如何用 PyTorch 的 nn.TransformerEncoder 复现 ViT 论文中的 Multi-Head Attention 与 Transformer Encoder【免费下载链接】pytorch-deep-learningMaterials for the Learn PyTorch for Deep Learning: Zero to Mastery course.项目地址: https://gitcode.com/GitHub_Trending/py/pytorch-deep-learning在 Learn PyTorch for Deep LearningZero to Mastery课程的 Paper Replicating 章节08_pytorch_paper_replicating.ipynb中ViT 论文的 Transformer Encoder 部分论文公式 2 的 Multi-Head Attention 块与公式 3 的 MLP 块会先手工写成自定义 block随后用 PyTorch 内置的torch.nn.TransformerEncoderLayer()与torch.nn.TransformerEncoder()复现出结构相同的编码器。这篇文章走的就是内置层这条路径从 ViT 论文 Table 1 / Table 3 的 ViT-Base 超参数出发构建单个编码器层、叠成 12 层并验证它与手工 block 在层组成、参数量和输入输出形状上一致。适用环境torch 1.12和torchvision 0.13。文档的练习部分明确写着练习要求torchvision v0.132022 年 7 月发布更早的版本可能可以运行但很可能出错解答 notebook 里的版本断言要求 torch 主版本 12、torchvision 0.13。1. 准备条件版本与依赖检查先确认 torch / torchvision 版本满足要求。下面这段断言来自解答 notebookextras/solutions/08_pytorch_paper_replicating_exercise_solutions.ipynbimport torch import torchvision assert int(torch.__version__.split(.)[1]) 12, 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__})注意一个副作用解答 notebook 原版在断言失败时会通过pip3 install -U --pre torch torchvision torchaudio自动安装 PyTorch nightly 版本会修改当前环境。在普通 Python 脚本里不需要照搬这段逻辑直接自行安装满足上述版本要求的稳定版即可。后续验证会用到torchinfo解答 notebook 在缺失时执行pip install -q torchinfo同样是安装第三方包的操作提前安装pip install torchinfo2. ViT 的 Transformer Encoder 由哪些部分组成文档把 ViT 论文 Figure 1 的 Transformer Encoder 拆成两个公式逐段复现这也是理解内置层参数映射的基础引号内为 ViT 论文原文文档中的引用公式 2MSA blockz_l MSA(LN(z_{l-1})) z_{l-1}。即 Multi-Head Attention 外包一层 LayerNorm再加残差连接输入加回输出。文档说明 MSA 用torch.nn.MultiheadAttention()、LN 用torch.nn.LayerNorm()实现。公式 3MLP blockz_l MLP(LN(z_l)) z_l。文档按论文 section 3.1 与 Appendix B.1 给出的 MLP 结构为layer norm - linear layer - non-linear layer - dropout - linear layer - dropout其中非线性是 GELUThe MLP contains two layers with a GELU non-linearitydropout 在每个 dense layer 之后应用qkv-projections 除外。论文还要求 Layernorm (LN) is appliedbeforeevery block, andresidual connections after every block这对应内置层的norm_firstTrue。ViT-Base 的超参数文档代码注释中的来源Hidden size D 768Table 1、Heads 12Table 1、MLP size 3072Table 1、dropout 0.1Table 3、Layers 12Table 1即要叠 12 个编码器 block。文档在 section 5.3 先用自定义类复现公式 2核心代码如下完整代码见主 notebook 08_pytorch_paper_replicating.ipynbclass MultiheadSelfAttentionBlock(nn.Module): Creates a multi-head self-attention block (MSA block for short). def __init__(self, embedding_dim:int768, # Hidden size D from Table 1 for ViT-Base num_heads:int12, # Heads from Table 1 for ViT-Base attn_dropout:float0): # doesnt look like the paper uses any dropout in MSABlocks super().__init__() # 3. Create the Norm layer (LN) self.layer_norm nn.LayerNorm(normalized_shapeembedding_dim) # 4. Create the Multi-Head Attention (MSA) layer self.multihead_attn nn.MultiheadAttention(embed_dimembedding_dim, num_headsnum_heads, dropoutattn_dropout, batch_firstTrue) # does our batch dimension come first? # 5. Create a forward() method to pass the data through the layers def forward(self, x): x self.layer_norm(x) attn_output, _ self.multihead_attn(queryx, # query embeddings keyx, # key embeddings valuex, # value embeddings need_weightsFalse) # do we need the weights or just the layer outputs? return attn_output两个要点nn.MultiheadAttention的 q、k、v 三个输入都是 Norm 层输出的同一份张量文档称之为 query、key、value 三重输入这个手工 block 不含残差连接残差是到 section 7.1 组装TransformerEncoderBlock时才加的x self.msa_block(x) xMLP block 同理。内置层路线正是把公式 2 公式 3 pre-LN 残差打包成一个TransformerEncoderLayer。3. 用 nn.TransformerEncoderLayer 构建单个编码器层torch.nn.TransformerEncoderLayer需要与 ViT-Base 的超参数一一对应映射关系如下取值均为主 notebook 代码中的注释内置层参数取值对应 ViT 论文内容d_model768Hidden size DTable 1ViT-Basenhead12HeadsTable 1ViT-Basedim_feedforward3072MLP sizeTable 1ViT-Basedropout0.1dense 层 dropoutTable 3ViT-BaseactivationgeluThe MLP contains two layers with a GELU non-linearitybatch_firstTruebatch 维度在前batch_firstTrue文档注释norm_firstTrueLayernorm is applied before every block构建代码来自解答 notebook 的练习 1 解法from torch import nn # Hyperparameters from Table 1 and Table 3 for ViT-Base transformer_encoder_layer nn.TransformerEncoderLayer(d_model768, nhead12, dim_feedforward3072, dropout0.1, activationgelu, batch_firstTrue, norm_firstTrue) transformer_encoder_layer文档示例输出实际值以你的 PyTorch 版本为准TransformerEncoderLayer( (self_attn): MultiheadAttention( (out_proj): NonDynamicallyQuantizableLinear(in_features768, out_features768, biasTrue) ) (linear1): Linear(in_features768, out_features3072, biasTrue) (dropout): Dropout(p0.1, inplaceFalse) (linear2): Linear(in_features3072, out_features768, biasTrue) (norm1): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (norm2): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (dropout1): Dropout(p0.1, inplaceFalse) (dropout2): Dropout(p0.1, inplaceFalse) )从这个结构可以看到self_attn就是MultiheadAttentionMSA 部分out_proj为 768→768 的线性层linear1/linear2是 768→3072→768 的两层 MLPnorm1/norm2是两个 LayerNorm(768)——与第 2 节公式 2、公式 3 手工 block 的层组成一致。4. 用 nn.TransformerEncoder 叠成 12 层ViT-Base 的 Table 1 有 12 个 Layers。文档说明可以用torch.nn.TransformerEncoder(encoder_layer, num_layers)完成堆叠其中encoder_layer用torch.nn.TransformerEncoderLayer()创建的目标单层num_layers要堆叠的 Transformer Encoder 层数。transformer_encoder nn.TransformerEncoder( encoder_layertransformer_encoder_layer, num_layers12)5. 验证内置编码器与手工 block 一致文档给出的判定标准section 7.2 原文内置层 summary 的输出结构因torch.nn.TransformerEncoderLayer()的构造方式与手工 block 略有不同但the layers it uses, number of parameters and input and output shapes are the same。可以分两步核对。核对结构与前向形状。文档在 7.1/7.2 中用torchinfo.summary()打印输入(1, 197, 768)batch_size, num_patches, embedding_dimension其中 197 196 个 patch 1 个 class token经过编码器 block 的 summary并直接对 block 做前向测试文档示例输出Input shape of MSA block: torch.Size([1, 197, 768]) Output shape MSA block: torch.Size([1, 197, 768])对应到内置层可以写from torchinfo import summary summary(modeltransformer_encoder_layer, input_size(1, 197, 768)) # (batch_size, num_patches, embedding_dimension) x torch.randn(1, 197, 768) # (batch_size, num_patches, embedding_dimension)文档 block 测试用的形状 print(fInput shape: {x.shape}) out transformer_encoder_layer(x) print(fOutput shape: {out.shape})前向测试的判定标准是输出形状与输入相同仍为(1, 197, 768)上面 5.3 节的文档示例输出即此结果。核对参数量。解答 notebook 对单层的 summary文档示例输出摘录TransformerEncoderLayer [32, 196, 768] 3,072 ├─LayerNorm: 1-6 [32, 196, 768] (recursive) ├─MultiheadAttention: 1-2 [32, 196, 768] 2,362,368 ├─Linear: 1-5 [32, 196, 3072] 2,362,368 ├─Linear: 1-10 [32, 196, 768] 2,360,064 ... Total params: 7,087,872 Trainable params: 7,087,872主 notebook 对手工TransformerEncoderBlock的 summary 同样给出输入输出(1, 197, 768)不变、参数量一致见 7.1 节附带的 summary 图images/08-vit-paper-summary-output-transformer-encoder.png。对照两边 summary 时按文档结论关注三点使用的层、参数量、输入输出形状是否相同而不要逐行比对 summary 的排列结构。6. 可选扩展作为完整 ViT 的编码器主干主 notebook 练习 1 的目标是用内置 PyTorch Transformer 层复现我们创建的 ViT 架构把自定义TransformerEncoderBlock()替换为nn.TransformerEncoderLayer()再用nn.TransformerEncoder()堆叠。解答 notebook 给出的完整代码如下其中PatchEmbedding类沿用主 notebook section 4.5 的实现nn.Conv2d打 patch nn.Flattenpermuteclass ViT(nn.Module): def __init__(self, img_size224, # from Table 3 num_channels3, patch_size16, embedding_dim768, # from Table 1 dropout0.1, mlp_size3072, # from Table 1 num_transformer_layers12, # from Table 1 num_heads12, # from Table 1 (number of multi-head self attention heads) num_classes1000): # generic number of classes (this can be adjusted) super().__init__() # Assert image size is divisible by patch size assert img_size % patch_size 0, Image size must be divisble by patch size. # 1. Create patch embedding self.patch_embedding PatchEmbedding(in_channelsnum_channels, patch_sizepatch_size, embedding_dimembedding_dim) # 2. Create class token self.class_token nn.Parameter(torch.randn(1, 1, embedding_dim), requires_gradTrue) # 3. Create positional embedding num_patches (img_size * img_size) // patch_size**2 # N HW/P^2 self.positional_embedding nn.Parameter(torch.randn(1, num_patches1, embedding_dim)) # 4. Create patch position embedding dropout self.embedding_dropout nn.Dropout(pdropout) # 5. Create stack Transformer Encoder layers (stacked single layers) self.transformer_encoder nn.TransformerEncoder(encoder_layernn.TransformerEncoderLayer(d_modelembedding_dim, nheadnum_heads, dim_feedforwardmlp_size, activationgelu, batch_firstTrue, norm_firstTrue), # Create a single Transformer Encoder Layer num_layersnum_transformer_layers) # Stack it N times # 7. Create MLP head self.mlp_head nn.Sequential( nn.LayerNorm(normalized_shapeembedding_dim), nn.Linear(in_featuresembedding_dim, out_featuresnum_classes) ) def forward(self, x): # Get some dimensions from x batch_size x.shape[0] # Create the patch embedding x self.patch_embedding(x) # First, expand the class token across the batch size class_token self.class_token.expand(batch_size, -1, -1) # -1 means infer the dimension # Prepend the class token to the patch embedding x torch.cat((class_token, x), dim1) # Add the positional embedding to patch embedding with class token x self.positional_embedding x # Dropout on patch positional embedding x self.embedding_dropout(x) # Pass embedding through Transformer Encoder stack x self.transformer_encoder(x) # Pass 0th index of x through MLP head x self.mlp_head(x[:, 0]) return x验证方式与解答 notebook 相同构造一张(1, 3, 224, 224)的随机图像张量前向一遍224 是 Table 3 的训练分辨率16 的 patch 打出 196 个 patch加上 class token 后序列长度为 197device cuda if torch.cuda.is_available() else cpu demo_img torch.randn(1, 3, 224, 224).to(device) print(demo_img.shape) # Create ViT vit ViT(num_classes3).to(device) # 3 是文档 pizza/steak/sushi 数据集的类别数按你的任务替换 out vit(demo_img) print(out.shape)解答 notebook 原文使用ViT(num_classeslen(class_names))其中class_names来自 pizza_steak_sushi 数据[pizza, steak, sushi]这里取其等价值 3 展示主 notebook 中 ViT 的默认值是num_classes1000ImageNet 默认。文档示例输出具体数值取决于随机初始化权重不必相同torch.Size([1, 3, 224, 224]) tensor([[-1.2707, -0.5487, 0.1726]], devicecuda:0, grad_fnAddmmBackward0)进一步可以用summary(modelViT(num_classes3), input_sizedemo_img.shape)核对整体结构。解答 notebook 的文档示例输出摘录显示输入(1, 3, 224, 224)经 PatchEmbedding 变为(1, 196, 768)加 class token 后(1, 197, 768)TransformerEncoder下挂着 12 个TransformerEncoderLayer每个 7,087,872 参数模型总参数 85,800,963num_classes3 时输出(1, 3)。7. 限制与文档注意事项summary 结构不必逐行相同文档明确说明内置层的 summary 输出因构造方式与手工 block 略有不同判定一致性的依据是使用的层、参数量、输入输出形状相同不要以 summary 文本完全一致作为成功条件。注意力投影的 dropout手工 MSA block 取attn_dropout0文档注释according to Appendix B.1, dropout isnt used after the qkv-projections文档示例中内置层只传了dropout0.1没有说明如何在内置层中单独控制注意力投影的 dropout如需严格对齐论文该细节这一点文档未覆盖。性能说法文档在 7.2 节指出条件为 2022 年 7 月、PyTorch 1.12内置torch.nn.TransformerEncoderLayer()在许多常见工作负载上可获得超过 2 倍的速度提升引用 PyTorch 官方博客并说明内置层更少出错、性能可能更好。这是文档给出的选型理由数值以上述时间条件为限。版本下限torch 1.12 / torchvision 0.13低于该版本很可能出错文档原文。手工 block 的完整代码、公式 1–4 的逐步映射插图和 ViT 整体训练流程见课程主 notebook 08_pytorch_paper_replicating.ipynb本文内置层路线的代码以练习 1 解答 extras/solutions/08_pytorch_paper_replicating_exercise_solutions.ipynb 为准可继续参考该 notebook 后续练习把 ViT 转成独立脚本、在 pizza/steak/sushi 数据上训练预训练 ViT 特征提取器等。【免费下载链接】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），仅供参考 如何用 PyTorch 的 nn.TransformerEncoder 复现 ViT 论文中的 Multi-Head Attention 与 Transformer Encoder？
如何用 PyTorch 的 nn.TransformerEncoder 复现 ViT 论文中的 Multi-Head Attention 与 Transformer Encoder【免费下载链接】pytorch-deep-learningMaterials for the Learn PyTorch for Deep Learning: Zero to Mastery course.项目地址: https://gitcode.com/GitHub_Trending/py/pytorch-deep-learning在 Learn PyTorch for Deep LearningZero to Mastery课程的 Paper Replicating 章节08_pytorch_paper_replicating.ipynb中ViT 论文的 Transformer Encoder 部分论文公式 2 的 Multi-Head Attention 块与公式 3 的 MLP 块会先手工写成自定义 block随后用 PyTorch 内置的torch.nn.TransformerEncoderLayer()与torch.nn.TransformerEncoder()复现出结构相同的编码器。这篇文章走的就是内置层这条路径从 ViT 论文 Table 1 / Table 3 的 ViT-Base 超参数出发构建单个编码器层、叠成 12 层并验证它与手工 block 在层组成、参数量和输入输出形状上一致。适用环境torch 1.12和torchvision 0.13。文档的练习部分明确写着练习要求torchvision v0.132022 年 7 月发布更早的版本可能可以运行但很可能出错解答 notebook 里的版本断言要求 torch 主版本 12、torchvision 0.13。1. 准备条件版本与依赖检查先确认 torch / torchvision 版本满足要求。下面这段断言来自解答 notebookextras/solutions/08_pytorch_paper_replicating_exercise_solutions.ipynbimport torch import torchvision assert int(torch.__version__.split(.)[1]) 12, 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__})注意一个副作用解答 notebook 原版在断言失败时会通过pip3 install -U --pre torch torchvision torchaudio自动安装 PyTorch nightly 版本会修改当前环境。在普通 Python 脚本里不需要照搬这段逻辑直接自行安装满足上述版本要求的稳定版即可。后续验证会用到torchinfo解答 notebook 在缺失时执行pip install -q torchinfo同样是安装第三方包的操作提前安装pip install torchinfo2. ViT 的 Transformer Encoder 由哪些部分组成文档把 ViT 论文 Figure 1 的 Transformer Encoder 拆成两个公式逐段复现这也是理解内置层参数映射的基础引号内为 ViT 论文原文文档中的引用公式 2MSA blockz_l MSA(LN(z_{l-1})) z_{l-1}。即 Multi-Head Attention 外包一层 LayerNorm再加残差连接输入加回输出。文档说明 MSA 用torch.nn.MultiheadAttention()、LN 用torch.nn.LayerNorm()实现。公式 3MLP blockz_l MLP(LN(z_l)) z_l。文档按论文 section 3.1 与 Appendix B.1 给出的 MLP 结构为layer norm - linear layer - non-linear layer - dropout - linear layer - dropout其中非线性是 GELUThe MLP contains two layers with a GELU non-linearitydropout 在每个 dense layer 之后应用qkv-projections 除外。论文还要求 Layernorm (LN) is appliedbeforeevery block, andresidual connections after every block这对应内置层的norm_firstTrue。ViT-Base 的超参数文档代码注释中的来源Hidden size D 768Table 1、Heads 12Table 1、MLP size 3072Table 1、dropout 0.1Table 3、Layers 12Table 1即要叠 12 个编码器 block。文档在 section 5.3 先用自定义类复现公式 2核心代码如下完整代码见主 notebook 08_pytorch_paper_replicating.ipynbclass MultiheadSelfAttentionBlock(nn.Module): Creates a multi-head self-attention block (MSA block for short). def __init__(self, embedding_dim:int768, # Hidden size D from Table 1 for ViT-Base num_heads:int12, # Heads from Table 1 for ViT-Base attn_dropout:float0): # doesnt look like the paper uses any dropout in MSABlocks super().__init__() # 3. Create the Norm layer (LN) self.layer_norm nn.LayerNorm(normalized_shapeembedding_dim) # 4. Create the Multi-Head Attention (MSA) layer self.multihead_attn nn.MultiheadAttention(embed_dimembedding_dim, num_headsnum_heads, dropoutattn_dropout, batch_firstTrue) # does our batch dimension come first? # 5. Create a forward() method to pass the data through the layers def forward(self, x): x self.layer_norm(x) attn_output, _ self.multihead_attn(queryx, # query embeddings keyx, # key embeddings valuex, # value embeddings need_weightsFalse) # do we need the weights or just the layer outputs? return attn_output两个要点nn.MultiheadAttention的 q、k、v 三个输入都是 Norm 层输出的同一份张量文档称之为 query、key、value 三重输入这个手工 block 不含残差连接残差是到 section 7.1 组装TransformerEncoderBlock时才加的x self.msa_block(x) xMLP block 同理。内置层路线正是把公式 2 公式 3 pre-LN 残差打包成一个TransformerEncoderLayer。3. 用 nn.TransformerEncoderLayer 构建单个编码器层torch.nn.TransformerEncoderLayer需要与 ViT-Base 的超参数一一对应映射关系如下取值均为主 notebook 代码中的注释内置层参数取值对应 ViT 论文内容d_model768Hidden size DTable 1ViT-Basenhead12HeadsTable 1ViT-Basedim_feedforward3072MLP sizeTable 1ViT-Basedropout0.1dense 层 dropoutTable 3ViT-BaseactivationgeluThe MLP contains two layers with a GELU non-linearitybatch_firstTruebatch 维度在前batch_firstTrue文档注释norm_firstTrueLayernorm is applied before every block构建代码来自解答 notebook 的练习 1 解法from torch import nn # Hyperparameters from Table 1 and Table 3 for ViT-Base transformer_encoder_layer nn.TransformerEncoderLayer(d_model768, nhead12, dim_feedforward3072, dropout0.1, activationgelu, batch_firstTrue, norm_firstTrue) transformer_encoder_layer文档示例输出实际值以你的 PyTorch 版本为准TransformerEncoderLayer( (self_attn): MultiheadAttention( (out_proj): NonDynamicallyQuantizableLinear(in_features768, out_features768, biasTrue) ) (linear1): Linear(in_features768, out_features3072, biasTrue) (dropout): Dropout(p0.1, inplaceFalse) (linear2): Linear(in_features3072, out_features768, biasTrue) (norm1): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (norm2): LayerNorm((768,), eps1e-05, elementwise_affineTrue) (dropout1): Dropout(p0.1, inplaceFalse) (dropout2): Dropout(p0.1, inplaceFalse) )从这个结构可以看到self_attn就是MultiheadAttentionMSA 部分out_proj为 768→768 的线性层linear1/linear2是 768→3072→768 的两层 MLPnorm1/norm2是两个 LayerNorm(768)——与第 2 节公式 2、公式 3 手工 block 的层组成一致。4. 用 nn.TransformerEncoder 叠成 12 层ViT-Base 的 Table 1 有 12 个 Layers。文档说明可以用torch.nn.TransformerEncoder(encoder_layer, num_layers)完成堆叠其中encoder_layer用torch.nn.TransformerEncoderLayer()创建的目标单层num_layers要堆叠的 Transformer Encoder 层数。transformer_encoder nn.TransformerEncoder( encoder_layertransformer_encoder_layer, num_layers12)5. 验证内置编码器与手工 block 一致文档给出的判定标准section 7.2 原文内置层 summary 的输出结构因torch.nn.TransformerEncoderLayer()的构造方式与手工 block 略有不同但the layers it uses, number of parameters and input and output shapes are the same。可以分两步核对。核对结构与前向形状。文档在 7.1/7.2 中用torchinfo.summary()打印输入(1, 197, 768)batch_size, num_patches, embedding_dimension其中 197 196 个 patch 1 个 class token经过编码器 block 的 summary并直接对 block 做前向测试文档示例输出Input shape of MSA block: torch.Size([1, 197, 768]) Output shape MSA block: torch.Size([1, 197, 768])对应到内置层可以写from torchinfo import summary summary(modeltransformer_encoder_layer, input_size(1, 197, 768)) # (batch_size, num_patches, embedding_dimension) x torch.randn(1, 197, 768) # (batch_size, num_patches, embedding_dimension)文档 block 测试用的形状 print(fInput shape: {x.shape}) out transformer_encoder_layer(x) print(fOutput shape: {out.shape})前向测试的判定标准是输出形状与输入相同仍为(1, 197, 768)上面 5.3 节的文档示例输出即此结果。核对参数量。解答 notebook 对单层的 summary文档示例输出摘录TransformerEncoderLayer [32, 196, 768] 3,072 ├─LayerNorm: 1-6 [32, 196, 768] (recursive) ├─MultiheadAttention: 1-2 [32, 196, 768] 2,362,368 ├─Linear: 1-5 [32, 196, 3072] 2,362,368 ├─Linear: 1-10 [32, 196, 768] 2,360,064 ... Total params: 7,087,872 Trainable params: 7,087,872主 notebook 对手工TransformerEncoderBlock的 summary 同样给出输入输出(1, 197, 768)不变、参数量一致见 7.1 节附带的 summary 图images/08-vit-paper-summary-output-transformer-encoder.png。对照两边 summary 时按文档结论关注三点使用的层、参数量、输入输出形状是否相同而不要逐行比对 summary 的排列结构。6. 可选扩展作为完整 ViT 的编码器主干主 notebook 练习 1 的目标是用内置 PyTorch Transformer 层复现我们创建的 ViT 架构把自定义TransformerEncoderBlock()替换为nn.TransformerEncoderLayer()再用nn.TransformerEncoder()堆叠。解答 notebook 给出的完整代码如下其中PatchEmbedding类沿用主 notebook section 4.5 的实现nn.Conv2d打 patch nn.Flattenpermuteclass ViT(nn.Module): def __init__(self, img_size224, # from Table 3 num_channels3, patch_size16, embedding_dim768, # from Table 1 dropout0.1, mlp_size3072, # from Table 1 num_transformer_layers12, # from Table 1 num_heads12, # from Table 1 (number of multi-head self attention heads) num_classes1000): # generic number of classes (this can be adjusted) super().__init__() # Assert image size is divisible by patch size assert img_size % patch_size 0, Image size must be divisble by patch size. # 1. Create patch embedding self.patch_embedding PatchEmbedding(in_channelsnum_channels, patch_sizepatch_size, embedding_dimembedding_dim) # 2. Create class token self.class_token nn.Parameter(torch.randn(1, 1, embedding_dim), requires_gradTrue) # 3. Create positional embedding num_patches (img_size * img_size) // patch_size**2 # N HW/P^2 self.positional_embedding nn.Parameter(torch.randn(1, num_patches1, embedding_dim)) # 4. Create patch position embedding dropout self.embedding_dropout nn.Dropout(pdropout) # 5. Create stack Transformer Encoder layers (stacked single layers) self.transformer_encoder nn.TransformerEncoder(encoder_layernn.TransformerEncoderLayer(d_modelembedding_dim, nheadnum_heads, dim_feedforwardmlp_size, activationgelu, batch_firstTrue, norm_firstTrue), # Create a single Transformer Encoder Layer num_layersnum_transformer_layers) # Stack it N times # 7. Create MLP head self.mlp_head nn.Sequential( nn.LayerNorm(normalized_shapeembedding_dim), nn.Linear(in_featuresembedding_dim, out_featuresnum_classes) ) def forward(self, x): # Get some dimensions from x batch_size x.shape[0] # Create the patch embedding x self.patch_embedding(x) # First, expand the class token across the batch size class_token self.class_token.expand(batch_size, -1, -1) # -1 means infer the dimension # Prepend the class token to the patch embedding x torch.cat((class_token, x), dim1) # Add the positional embedding to patch embedding with class token x self.positional_embedding x # Dropout on patch positional embedding x self.embedding_dropout(x) # Pass embedding through Transformer Encoder stack x self.transformer_encoder(x) # Pass 0th index of x through MLP head x self.mlp_head(x[:, 0]) return x验证方式与解答 notebook 相同构造一张(1, 3, 224, 224)的随机图像张量前向一遍224 是 Table 3 的训练分辨率16 的 patch 打出 196 个 patch加上 class token 后序列长度为 197device cuda if torch.cuda.is_available() else cpu demo_img torch.randn(1, 3, 224, 224).to(device) print(demo_img.shape) # Create ViT vit ViT(num_classes3).to(device) # 3 是文档 pizza/steak/sushi 数据集的类别数按你的任务替换 out vit(demo_img) print(out.shape)解答 notebook 原文使用ViT(num_classeslen(class_names))其中class_names来自 pizza_steak_sushi 数据[pizza, steak, sushi]这里取其等价值 3 展示主 notebook 中 ViT 的默认值是num_classes1000ImageNet 默认。文档示例输出具体数值取决于随机初始化权重不必相同torch.Size([1, 3, 224, 224]) tensor([[-1.2707, -0.5487, 0.1726]], devicecuda:0, grad_fnAddmmBackward0)进一步可以用summary(modelViT(num_classes3), input_sizedemo_img.shape)核对整体结构。解答 notebook 的文档示例输出摘录显示输入(1, 3, 224, 224)经 PatchEmbedding 变为(1, 196, 768)加 class token 后(1, 197, 768)TransformerEncoder下挂着 12 个TransformerEncoderLayer每个 7,087,872 参数模型总参数 85,800,963num_classes3 时输出(1, 3)。7. 限制与文档注意事项summary 结构不必逐行相同文档明确说明内置层的 summary 输出因构造方式与手工 block 略有不同判定一致性的依据是使用的层、参数量、输入输出形状相同不要以 summary 文本完全一致作为成功条件。注意力投影的 dropout手工 MSA block 取attn_dropout0文档注释according to Appendix B.1, dropout isnt used after the qkv-projections文档示例中内置层只传了dropout0.1没有说明如何在内置层中单独控制注意力投影的 dropout如需严格对齐论文该细节这一点文档未覆盖。性能说法文档在 7.2 节指出条件为 2022 年 7 月、PyTorch 1.12内置torch.nn.TransformerEncoderLayer()在许多常见工作负载上可获得超过 2 倍的速度提升引用 PyTorch 官方博客并说明内置层更少出错、性能可能更好。这是文档给出的选型理由数值以上述时间条件为限。版本下限torch 1.12 / torchvision 0.13低于该版本很可能出错文档原文。手工 block 的完整代码、公式 1–4 的逐步映射插图和 ViT 整体训练流程见课程主 notebook 08_pytorch_paper_replicating.ipynb本文内置层路线的代码以练习 1 解答 extras/solutions/08_pytorch_paper_replicating_exercise_solutions.ipynb 为准可继续参考该 notebook 后续练习把 ViT 转成独立脚本、在 pizza/steak/sushi 数据上训练预训练 ViT 特征提取器等。【免费下载链接】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），仅供参考