Vision Transformer（vit）的Encoder层结构

图解：

代码：

class Block(nn.Module):
    def __init__(self,
                 dim,                       #输入特征的维度
                 num_heads,                 #多头注意力机制中的注意力头的数量
                 mlp_ratio=4.,              #MLP模块中隐藏层大小与输入维度的比例，默认是 4
                 qkv_bias=False,            #qkv的偏置默认false
                 qk_scale=None,             #缩放因子默认none
                 drop_ratio=0.,             #注意力和MLP层的 dropout 概率，默认是0。
                 attn_drop_ratio=0.,        #专门用于注意力层的 dropout 概率
                 drop_path_ratio=0.,        #用于droppath的dropout 概率
                 act_layer=nn.GELU,         #用于 MLP 中的激活函数，默认是 nn.GELU
                 norm_layer=nn.LayerNorm):  #规范化层，默认是 nn.LayerNorm
        super(Block, self).__init__()
        self.norm1 = norm_layer(dim)
#这是输入经过注意力模块前的规范化层，默认是 nn.LayerNorm
        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)
#初始化了注意力模块，实现了多头注意力机制
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()
#条件性的加入了DropPath，如果 drop_path_ratio大于0，就会应用 DropPath，否则不做任何操作。
        self.norm2 = norm_layer(dim)
#这是 MLP 模块之前的规范化层。
        mlp_hidden_dim = int(dim * mlp_ratio)
#计算 MLP 隐藏层的维度，它等于输入维度乘以 mlp_ratio
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,act_layer=act_layer, drop=drop_ratio)
#初始化了 MLP 模块

    def forward(self, x):
        x = x + self.drop_path(self.attn(self.norm1(x)))
#先规范，再注意力，再droppath，再残差连接
        x = x + self.drop_path(self.mlp(self.norm2(x)))
#先规范，再注意力，再droppath，再残差连接
        return x

注意：

Multi-Head Self-Attention：Vision Transformer（vit）的Multi-Head Self-Attention（多头注意力机制）结构-优快云博客

mlp结构：Vision Transformer（vit）的MLP模块-优快云博客

mlp_hidden_dim = int(dim * mlp_ratio)
#计算 MLP 隐藏层的维度，它等于输入维度乘以 mlp_ratio

mlp有一个隐藏层一个输出层隐藏层扩大mlp_ratio输出层恢复。