Vision transformer复现

Vision transformer复现

深度学习模型之——Stochastic depth（随机深度）

伯努利分布方差_概率分布，先懂这6个

Paddle 2.0：Vision Transformer 模型的构建

思路是 把一张图片划分成num_patches个小图片，然后对每一个图片做卷积，相当于将每个小图片生成embed（特征数量）个数，用这些数表示小图片，剩下的就是正常的transformer了

import numpy as np
import paddle
import paddle.nn as nn
from paddle.nn.initializer import TruncatedNormal, Constant

# 参数初始化配置
trunc_normal_ = TruncatedNormal(std=.02)
zeros_ = Constant(value=0.)
ones_ = Constant(value=1.)

# x[int] -> tuple(x, x)
def to_2tuple(x):
    return tuple([x] * 2)

# 独立层，即什么操作都没有的网络层
class Identity(nn.Layer):
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, input):
        return input

class PatchEmbed(nn.Layer):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        num_patches = (img_size[1] // patch_size[1]) * \
            (img_size[0] // patch_size[0])
        # 将图片划分成 num_patches 个大小为 patch_size**2的图像
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches

        self.proj = nn.Conv2D(in_chans, embed_dim,
                              kernel_size=patch_size, stride=patch_size)
        # 然后用卷积将 每个大小为 patch_size**2的图像，转化为具体数字
        # 这里会转换成embed_dim个数字，相当于一个维度为embed_dim的向量
        # 运用这个向量就可以表示一个图片，相当于一个单词的特征向量
        # 输出为 B，embed，H/patch_size,W/patch_size


    def forward(self, x):
        B, C, H, W = x.shape
        # 分块线性变换 + 向量展平 + 维度转置
        """
        B:这一组batch数据有多少个
        C：通道数
        """
        x = self.proj(x).flatten(2).transpose((0, 2, 1))
        # 输出为：B,num_patches,embed
        return x

class Attention(nn.Layer):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, C = x.shape
        # 线性变换
        qkv = self.qkv(x).reshape((B, N, 3, self.num_heads, C //
                                   self.num_heads)).transpose((2, 0, 3, 1, 4))
        # 分割 query key value
        q, k, v = qkv[0], qkv[1], qkv[2]
        # Matmul + Scale
        attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
        # SoftMax
        attn = nn.functional.softmax(attn, axis=-1)
        # Attention Dropout
        attn = self.attn_drop(attn)
        # Matmul
        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((B, N, C))
        # 线性变换
        x = self.proj(x)
        # Linear Dropout
        x = self.proj_drop(x)
        return x

class Mlp(nn.Layer):
    # transformer完后进行全连接层
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        # 输入层：线性变换
        x = self.fc1(x)
        # 应用激活函数
        x = self.act(x)
        # Dropout
        x = self.drop(x)
        # 输出层：线性变换
        x = self.fc2(x)
        # Dropout
        x = self.drop(x)
        return x

def drop_path(x,drop_p=0.,training=False):
    # 元组第一个维数是batch
    # 操作结束后相当于是有 batch*drop_p的，每个batch对应的张量，的全部元素是0
    # 相当于丢去
    if drop_p==0 or not training:
        return x
    keep_p=paddle.to_tensor(1-drop_p)
    shape=(x.shape[0],)+(1,)*(x.ndim-1)
    # 这里产生的操作是：生成一个元组，来表示shape，比如(5,4,3,2)
    # 元组第一个是batch，后面要和x相同
    """
    表示batch个张量中，对于每一个要被处理的张量
    后面都会有一个是否drop掉
    """
    random_tensor=keep_p+paddle.rand(shape,dtype=x.dtype)
    # 随机会生成一个张量，每个元素大小是0到1之间,均匀分布
    # 加上那个概率，就会让 这个概率的个数的数 是1以上，其他是0到1
    # 不能用randn，那样会生成正态分布的向量
    random_tensor=paddle.floor(random_tensor)
    # 把向量里面每个元素转换成0,1
    output=x.divide(keep_p)*random_tensor
    return output

class DropPath(nn.Layer):
    def __init__(self,drop_p=None):
        super(DropPath, self).__init__()
        self.drop_p=drop_p
    def forward(self, x):
        return drop_path(x,self.drop_p,self.training)
class Block(nn.Layer):
    def __init__(self, dim, num_heads, mlp_ratio=4.,
                 qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU,
                 norm_layer='nn.LayerNorm', epsilon=1e-5):
        """
        :param epsilon: 表示极小值，用于处理类似于log(x+epsilon),防止x为0
        """
        self.norm1=eval(norm_layer,epsilon=epsilon)
        self.attn=Attention(
            dim,num_heads=num_heads,qkv_bias=qkv_bias,
            qk_scale=qk_scale, attn_drop=attn_drop,
            proj_drop=drop
        )
        self.drop_path=DropPath(drop_path) if drop_path>0. else Identity()
        # Identity是空层
        self.norm2 = eval(norm_layer)(dim, epsilon=epsilon)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
                       act_layer=act_layer, drop=drop)
    def forward(self, x):
        x=x+self.drop_path(self.attn(self.norm1(x)))
        x=x+self.drop_path(self.mlp(self.norm2(x)))

class VisionTransformer(nn.Layer):
    def __init__(self, img_size=224, patch_size=16, in_chans=3,
                 class_dim=1000, embed_dim=768, depth=12,
                 num_heads=12, mlp_ratio=4, qkv_bias=False,
                 qk_scale=None, drop_rate=0., attn_drop_rate=0.,
                 drop_path_rate=0., norm_layer='nn.LayerNorm',
                 epsilon=1e-5, **args):
        super.__init__()
        self.class_dim=class_dim
        self.num_features=self.embed_dim=embed_dim
        self.patch_embed=PatchEmbed(
            img_size=img_size,patch_size=patch_size
            ,in_chans=in_chans,embed_dim=embed_dim
        )
        num_patches=self.patch_embed.num_patches
        self.pos_embed=self.create_parameter(
            shape=(1,num_patches+1,embed_dim),
            default_initializer=zeros_
        )
        self.add_parameter('pos_embed',self.pos_embed)
        self.cls_token=self.create_parameter(
            shape=(1,1,embed_dim),default_initializer=zeros_
        )
        self.add_parameter("cls_token",self.cls_token)

        self.pos_drop=nn.Dropout(p=drop_rate)
        # 初始时位置信息的dropout

        """
        这里运用了Stochastic depth
        """
        dpr=[x for x in paddle.linspace(0,drop_path_rate,depth)]
        self.blocks=nn.LayerList([
            Block(dim=embed_dim, num_heads=num_heads,
                  mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
                drop=drop_rate, attn_drop=attn_drop_rate,
                  drop_path=dpr[i], norm_layer=norm_layer, epsilon=epsilon)
        for i in range(depth)])
        self.norm=eval(norm_layer)(embed_dim,epsilon=epsilon)
        self.head=nn.Linear(embed_dim,class_dim) if class_dim>0 else Identity()
        trunc_normal_(self.pos_embed)
        trunc_normal_(self.cls_token)
        self.apply(self._init_weight)
    def _init_weight(self,m):
        if isinstance(m,nn.Linear):
            trunc_normal_(m.weight)
            if m.bias is not None:
                zeros_(m.bias)
        elif isinstance(m,nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)
    def forward_features(self,x):
        B=x.shape[0]
        x=self.patch_embed(x)
        # x是B,num_patch,embed
        cls_tokens=self.cls_token.expand((B,-1,-1))
        # cls_tokens是B,1,embed
        x=paddle.concat((cls_tokens,x),axis=1)
        # x是B,num_patch+1,embed
        x=x+self.pos_embed
        # self.pos_embed是1,num_patch+1,embed
        x=self.pos_drop(x)
        
        for blk in self.blocks:
            x=blk(x)
        x=self.norm(x)
        return x[:,0]
    def forward(self,x):
        x=self.forward_features(x)
        x=self.head(x)
        return x