（pytorch进阶之路四）Vision Transformer

大致思想

Vision Transformer 用的是Encoder only类型，主要用的就是Transformer Encoder模块

VIT的想法就是将Transformer应用到图像识别上去

但是直接应用有个问题，NLP是单词为单位，一句话的词数还是比较少的，但是图片的基本单位是一个个像素点，数量多得多，导致计算量会非常的大。

图片的一个像素点不包含多少信息量，对于图片它的信息量主要还是聚焦在一小块区域上。

直接的想法就是将很多的像素点组成块，图像分成很多个块，将图像块当作一个token送入到transformer中

块有两种角度理解，一是用DNN角度，将图片切割成很多个块，再将这很多个块经过仿射变换变成向量。

另一个角度是CNN角度，块用卷积操作，stride = kernel size，再将输出卷积图拉直变成一个个token向量。

为了做分类任务，VIT借鉴了bert中class token embedding这样一个占位符

在VIT中同样用了position embedding，经对比还是使用可训练的一维embedding好一点

最后用class token这个状态量做分类任务

将图片分成固定大小的块，用线性网络得到嵌入式表征，再加入位置编码，将一系列的向量送入到标准transformer中。

图片大小可能会变换，但是块的大小是不会变的，反应的只是序列长度大一些。

在encoder的最后一层，取出新加的输出状态，经过MLP得到类别的概率分布，使用交叉熵算出分类loss

论文地址

https://openreview.net/pdf?id=YicbFdNTTy

Patch embedding

第一步将图片变成embedding，两种方式变embedding，一是DNN，二是CNN

首先图片分块，其实就是卷积中的滑动操作，我们可以直接使用functional中的Unfold函数，直接拿出所有卷积的区域，刚好就是image to patch的过程

一个patch的大小 = ph * pw * ic

Unfold操作结果shape = [bs, patch大小, patch数量]

转置一下Unfold操作结果，再与DNN权重相乘，就得到了embedding

DNN embedding：

import torch.nn.functional as F
import torch as t

def image2emb_naive(image, patch_size, weight):
    """
    image shape: [bs,c,h,w]
    weight: DNN weight
    """
    patch = F.unfold(image, kernel_size=patch_size, stride=patch_size).transpose(-1, -2)
    patch_embedding = patch @ weight

    print(patch.shape, "# patch.shape")
    print(patch_embedding.shape, "# patch_embedding.shape")

    return patch_embedding


def test_image2emb_naive():
    bs, ic, ih, iw = 1, 3, 8, 8
    patch_size = 4
    model_dim = 8
    patch_depth = patch_size * patch_size * ic
    image = t.randn(bs, ic, ih, iw)
    weight = t.randn(patch_depth, model_dim)
    image2emb_naive(image, patch_size, weight)


test_image2emb_naive()

CNN卷积构造embedding最关键的是构造好kernel，model_dim其实就是输出通道数oc，

kernel shape = [oc, ic, kh, kw] = [model_dim, ic, patch_size, patch_size]

stride = patch_size

def image2emb_conv(image, kernel, stride):
    # conv_output: [bs, oc, oh, ow]
    conv_output = F.conv2d(image, kernel, stride=stride)
    bs, oc, oh, ow = conv_output.shape
    embedding = conv_output.reshape((bs, oc, oh*ow)).transpose(-1, -2)
    return embedding


def test_image2emb_conv():
    bs, ic, ih, iw = 1, 3, 8, 8
    patch_size = 4
    model_dim = 8
    image = t.randn(bs, ic, ih, iw)
    kernel = t.randn(model_dim, ic, patch_size, patch_size)
    patch_embedding = image2emb_conv(image, kernel, stride=patch_size)
    print(patch_embedding.shape)


test_image2emb_conv()

Class token

在序列前面添加CLS token，用来做分类任务

def append_cls_token(patch_embedding):
    bs, _, model_dim = patch_embedding.shape
    cls_token_embedding = t.randn(bs, 1, model_dim, requires_grad=True)
    # 把cls放到第一个位置上
    token_embedding = t.cat([cls_token_embedding, patch_embedding], dim=1)
    return token_embedding

Position embedding

这里的position embedding的思想类似word embedding，用一个table做embbeding

tile函数的作用就是复制，拷贝table和bs一样多份

def append_position_embedding(max_num_token, token_embedding):
    bs, seq_len, model_dim = token_embedding.shape
    # shape = [vocab_size, model_dim]
    position_embedding_table = t.randn(max_num_token, model_dim, requires_grad=True)
    position_embedding = t.tile(position_embedding_table[:seq_len], [bs, 1, 1])
    token_embedding += position_embedding
    return token_embedding

Encoder

直接调用nn API

def pass_embedding_to_encoder(token_embedding, ):
    bs, seq_len, model_dim = token_embedding.shape
    encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=8)
    encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
    encoder_output = encoder(token_embedding)
    return encoder_output

Classification mlp

最后是取出CLS token做分类任务

def do_classification(encoder_output, num_class, model_dim, label):
    # label = t.randint(10,(bs,))
    cls_token_output = encoder_output[:, 0, :]
    linear_layer = nn.Linear(model_dim, num_class)
    logits = linear_layer(cls_token_output)
    loss_fn = nn.CrossEntropyLoss()
    loss = loss_fn(logits, label)
    return loss

完整代码

import torch.nn.functional as F
import torch as t
import torch.nn as nn

def image2emb_conv(image, kernel, stride):
    # conv_output: [bs, oc, oh, ow]
    conv_output = F.conv2d(image, kernel, stride=stride)
    bs, oc, oh, ow = conv_output.shape
    embedding = conv_output.reshape((bs, oc, oh*ow)).transpose(-1, -2)
    return embedding

def append_cls_token(patch_embedding):
    bs, _, model_dim = patch_embedding.shape
    cls_token_embedding = t.randn(bs, 1, model_dim, requires_grad=True)
    # 把cls放到第一个位置上
    token_embedding = t.cat([cls_token_embedding, patch_embedding], dim=1)
    return token_embedding


def append_position_embedding(max_num_token, token_embedding):
    """
    max_num_token:序列最大长度
    """
    bs, seq_len, model_dim = token_embedding.shape
    # shape = [vocab_size, model_dim]
    position_embedding_table = t.randn(max_num_token, model_dim, requires_grad=True)
    position_embedding = t.tile(position_embedding_table[:seq_len], [bs, 1, 1])
    token_embedding += position_embedding
    return token_embedding


def pass_embedding_to_encoder(token_embedding):
    bs, seq_len, model_dim = token_embedding.shape
    encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=8)
    encoder = nn.TransformerEncoder(encoder_layer, num_layers=6)
    encoder_output = encoder(token_embedding)
    return encoder_output


def do_classification(encoder_output, num_class, model_dim, label):
    # label = t.randint(10,(bs,))
    cls_token_output = encoder_output[:, 0, :]
    linear_layer = nn.Linear(model_dim, num_class)
    logits = linear_layer(cls_token_output)
    loss_fn = nn.CrossEntropyLoss()
    loss = loss_fn(logits, label)
    return loss

def test_full():
    bs, ic, ih, iw = 1, 3, 8, 8
    patch_size = 4
    model_dim = 8
    max_num_token = 16
    num_class = 10
    label = t.randint(10,(bs,))
    image = t.randn(bs, ic, ih, iw)
    kernel = t.randn(model_dim, ic, patch_size, patch_size)
    patch_embedding = image2emb_conv(image, kernel, stride=patch_size)
    token_embedding = append_position_embedding(max_num_token ,append_cls_token(patch_embedding))
    encoder_output = pass_embedding_to_encoder(token_embedding)
    loss = do_classification(encoder_output, num_class, model_dim, label)
    print(loss)


test_full()