简介
本文的目的是通过实际代码编写来实现ViT模型，进一步加对ViT模型的理解，如果还不知道ViT模型的话，可以先看下博客了解一下ViT的整体结构。
本文整体是对Implementing Vision Transformer (ViT) in PyTorch 的翻译，但是也加上了一些自己的注解。如果读者更习惯看英文版，建议直接去看原文。
ViT模型整体结构
按照惯例，先放上模型的架构图，如下：
ViT模型
patch。接着这些patch被送入一个全连接层得到embeddings，然后在embeddings前前加上一个特殊的cls token。然后给所有的embedding加上位置信息编码positional encoding。接着这个整体被送入Transformer Encoder，然后取cls token的输出特征送入MLP Head去做分类，总体流程就是这样。
代码的整体结构跟ViT模型的结构类似，大体可以分为以下几个部分：
我们将以自底向上的方式来逐步实现ViT模型。
Data
首先需要导入相关的依赖库，如下：
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torch import nn
from torch import Tensor
from PIL import Image
from torchvision.transforms import Compose, Resize, ToTensor
from einops import rearrange, reduce, repeat
from einops.layers.torch import Rearrange, Reduce
from torchsummary import summary
首先我们需要打开一张图片，如下：
img = Image.open('./cat.jpg')
fig = plt.figure()
plt.imshow(img)
结果：
Cat
接着我们需要对图片进行预处理，主要是包含
resize、向量化等操作，代码如下：
# resize to imagenet size
transform = Compose([Resize((224, 224)), ToTensor()])
x = transform(img)
x = x.unsqueeze(0) # 主要是为了添加batch这个维度
x.shape
输出
Pathches Embeddings
根据ViT模型结构，第一步是需要将图片划分为多个Patches，并且将其铺平。如下图：
einops库来简化代码编写，如下：
patch_size = 16 # 16 pixels
pathes = rearrange(x, 'b c (h s1) (w s2) -> b (h w) (s1 s2 c)', s1=patch_size, s2=patch_size)
pathes.shape
结果
关于
einops库的使用可以参考doc。这里解释一下这个结果[1,196,768]是怎么来的。我们知道原始图片向量x的大小为[1,3,224,224]，当我们使用16x16大小的patch对其进行分割的时候，一共可以划分为224x224/16/16 = 196个patches，其次每个patch大小为16x16x3=768，故大小为[1,196,768]。
接着我们需要将这些patches通过一个线性映射层。
PatchEmbedding的类来使代码更加整洁：
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
self.patch_size = patch_size
super().__init__()
self.projection = nn.Sequential(
# break-down the image in s1 x s2 patches and flat them
Rearrange('b c (h s1) (w s2) -> b (h w) (s1 s2 c)', s1=patch_size, s2=patch_size),
# 注意这里的隐层大小设置的也是768，可以配置
nn.Linear(patch_size * patch_size * in_channels, emb_size)
)
def forward(self, x: Tensor) -> Tensor:
x = self.projection(x)
return x
PatchEmbedding()(x).shape
Conv2d层来实现相同的功能。这是通过设置卷积核大小和步长均为patch_size来实现的。直观上来看，卷积操作是分别应用在每个patch上的。所以，我们可以先应用一个卷积层，然后再对结果进行铺平，改进如下：
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
self.patch_size = patch_size
super().__init__()
self.projection = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
# 将卷积操作后的patch铺平
Rearrange('b e h w -> b (h w) e'),
)
def forward(self, x: Tensor) -> Tensor:
x = self.projection(x)
return x
PatchEmbedding()(x).shape
CLS Token
下一步是对映射后的patches添加上cls token以及位置编码信息。cls token是一个随机初始化的torch Parameter对象，在forward方法中它需要被拷贝b次(b是batch的数量)，然后使用torch.cat函数添加到patch前面。
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768):
self.patch_size = patch_size
super().__init__()
self.proj = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
Rearrange('b e (h) (w) -> b (h w) e'),
)
# 生成一个维度为emb_size的向量当做cls_token
self.cls_token = nn.Parameter(torch.randn(1, 1, emb_size))
def forward(self, x: Tensor) -> Tensor:
b, _, _, _ = x.shape # 单独先将batch缓存起来
x = self.proj(x) # 进行卷积操作
# 将cls_token 扩展b次
cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
print(cls_tokens.shape)
print(x.shape)
# prepend the cls token to the input on the dim 1
x = torch.cat([cls_tokens, x], dim=1)
return x
PatchEmbedding()(x).shape
Position Embedding
目前为止，模型还对patches的在图像中的原始位置一无所知。我们需要传递给模型这些空间上的信息。可以有很多种方法来实现这个功能，在ViT中，我们让模型自己去学习这个。位置编码信息只是一个形状为[N_PATCHES+1(token)m EMBED_SIZE]的张量，它直接加到映射后的patches上。
class PatchEmbedding(nn.Module):
def __init__(self, in_channels: int = 3, patch_size: int = 16, emb_size: int = 768, img_size: int = 224):
self.patch_size = patch_size
super().__init__()
self.projection = nn.Sequential(
# using a conv layer instead of a linear one -> performance gains
nn.Conv2d(in_channels, emb_size, kernel_size=patch_size, stride=patch_size),
Rearrange('b e (h) (w) -> b (h w) e'),
)
self.cls_token = nn.Parameter(torch.randn(1,1, emb_size))
# 位置编码信息，一共有(img_size // patch_size)**2 + 1(cls token)个位置向量
self.positions = nn.Parameter(torch.randn((img_size // patch_size)**2 + 1, emb_size))
def forward(self, x: Tensor) -> Tensor:
b, _, _, _ = x.shape
x = self.projection(x)
cls_tokens = repeat(self.cls_token, '() n e -> b n e', b=b)
# prepend the cls token to the input
x = torch.cat([cls_tokens, x], dim=1)
# add position embedding
print(x.shape, self.positions.shape)
x += self.positions
return x
PatchEmbedding()(x).shape
我们首先定义位置embedding向量，然后在forward函数中将其加到线性映射后的patches向量上去。
Transformer
现在我们来实现Transformer模块。ViT模型中只使用了Transformer的Encoder部分，其整体架构如下：
Encoder架构
接下来依次实现。
Attention
attention部分有三个输入，分别是queries,keys,values矩阵，首先使用queries,keys矩阵去计算注意力矩阵，然后与values矩阵相乘，得到对应的输出。在下图中，multi-head注意力机制表示将输入划分成n份，然后将计算分到n个head上去。
Attentionn
nn.MultiAttention模块或者自己实现一个，为了完整起见，我将完整的MultiAttention代码贴出来：
class MultiHeadAttention(nn.Module):
def __init__(self, emb_size: int = 512, num_heads: int = 8, dropout: float = 0):
super().__init__()
self.emb_size = emb_size
self.num_heads = num_heads
self.keys = nn.Linear(emb_size, emb_size)
self.queries = nn.Linear(emb_size, emb_size)
self.values = nn.Linear(emb_size, emb_size)
self.att_drop = nn.Dropout(dropout)
self.projection = nn.Linear(emb_size, emb_size)
def forward(self, x : Tensor, mask: Tensor = None) -> Tensor:
# split keys, queries and values in num_heads
queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.num_heads)
keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.num_heads)
values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.num_heads)
# sum up over the last axis
energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len
if mask is not None:
fill_value = torch.finfo(torch.float32).min
energy.mask_fill(~mask, fill_value)
scaling = self.emb_size ** (1/2)
att = F.softmax(energy, dim=-1) / scaling
att = self.att_drop(att)
# sum up over the third axis
out = torch.einsum('bhal, bhlv -> bhav ', att, values)
out = rearrange(out, "b h n d -> b n (h d)")
out = self.projection(out)
return out
接下来，一步一步分析下上述代码。我们定义了4个全连接层，分别用于queries,keys,values,以及最后的线性映射层。关于这块更加详细的内容可以阅读The Illustrated Transformer 。主要的思想是使用queries和keys之间的乘积来计算输入序列中的每一个patch与剩余patch之间的匹配程度。然后使用这个匹配程度(数值)去对对应的values做缩放，再累加起来作为Encoder的输出。
forward方法将上一层的输出作为输入，使用三个线性映射层分别得到queries,keys,values。因为我们要实现multi-head注意力机制，我们需要将输出重排成多个head的形式。这一步是使用einops库的rearrange函数来完成的。
Queries,keys,values的形状是一样的，为了简便起见，它们都是基于同一个输入x。
queries = rearrange(self.queries(x), "b n (h d) -> b h n d", h=self.n_heads)
keys = rearrange(self.keys(x), "b n (h d) -> b h n d", h=self.n_heads)
values = rearrange(self.values(x), "b n (h d) -> b h n d", h=self.n_heads)
经过rearrange操作之后，Queries,keys,values的形状大小为[BATCH, HEADS, SEQUENCE_LEN, EMBEDDING_SIZE]. 然后我们将多个head拼接在一起就得到了最终的输出。
注意： 为了加快计算，我们可以使用单个矩阵一次性计算出Queries,keys,values。
改进后的代码如下：
class MultiHeadAttention(nn.Module):
def __init__(self, emb_size: int = 768, num_heads: int = 8, dropout: float = 0):
super().__init__()
self.emb_size = emb_size
self.num_heads = num_heads
# fuse the queries, keys and values in one matrix
self.qkv = nn.Linear(emb_size, emb_size * 3)
self.att_drop = nn.Dropout(dropout)
self.projection = nn.Linear(emb_size, emb_size)
def forward(self, x : Tensor, mask: Tensor = None) -> Tensor:
# split keys, queries and values in num_heads
print("1qkv's shape: ", self.qkv(x).shape)
qkv = rearrange(self.qkv(x), "b n (h d qkv) -> (qkv) b h n d", h=self.num_heads, qkv=3)
print("2qkv's shape: ", qkv.shape)
queries, keys, values = qkv[0], qkv[1], qkv[2]
print("queries's shape: ", queries.shape)
print("keys's shape: ", keys.shape)
print("values's shape: ", values.shape)
# sum up over the last axis
energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len
print("energy's shape: ", energy.shape)
if mask is not None:
fill_value = torch.finfo(torch.float32).min
energy.mask_fill(~mask, fill_value)
scaling = self.emb_size ** (1/2)
print("scaling: ", scaling)
att = F.softmax(energy, dim=-1) / scaling
print("att1' shape: ", att.shape)
att = self.att_drop(att)
print("att2' shape: ", att.shape)
# sum up over the third axis
out = torch.einsum('bhal, bhlv -> bhav ', att, values)
print("out1's shape: ", out.shape)
out = rearrange(out, "b h n d -> b n (h d)")
print("out2's shape: ", out.shape)
out = self.projection(out)
print("out3's shape: ", out.shape)
return out
patches_embedded = PatchEmbedding()(x)
print("patches_embedding's shape: ", patches_embedded.shape)
MultiHeadAttention()(patches_embedded).shape
output
Residuals
Transformer模块也包含了残差连接，如下图：
Skip Connection
我们可以单独封装一个处理残差连接的类如下：
class ResidualAdd(nn.Module):
def __init__(self, fn):
super().__init__()
self.fn = fn
def forward(self, x, **kwargs):
res = x
x = self.fn(x, **kwargs)
x += res
return x
接着attention层的输出首先通过BN层，紧跟其后一个全连接层，全连接层中采用了一个expansion因子来对输入进行上采样。同样这里也采用了类似resnet的残差连接的方式。如下图：
FFN
class FeedForwardBlock(nn.Sequential):
def __init__(self, emb_size: int, expansion: int = 4, drop_p: float = 0.):
super().__init__(
nn.Linear(emb_size, expansion * emb_size),
nn.GELU(),
nn.Dropout(drop_p),
nn.Linear(expansion * emb_size, emb_size),
)
作者说，不知道为什么，很少看见人们直接继承
nn.Sequential类，这样就可以避免重写forward方法了。
译者著，确实~又学到了一招。
最终，我们可以创建一个完整的Transformer Encoder 块了。
Encoder
ResidualAdd类，我们可以很优雅地定义出Transformer Encoder Block，如下：
class TransformerEncoderBlock(nn.Sequential):
def __init__(self,
emb_size: int = 768,
drop_p: float = 0.,
forward_expansion: int = 4,
forward_drop_p: float = 0.,
** kwargs):
super().__init__(
ResidualAdd(nn.Sequential(
nn.LayerNorm(emb_size),
MultiHeadAttention(emb_size, **kwargs),
nn.Dropout(drop_p)
)),
ResidualAdd(nn.Sequential(
nn.LayerNorm(emb_size),
FeedForwardBlock(
emb_size, expansion=forward_expansion, drop_p=forward_drop_p),
nn.Dropout(drop_p)
)
))
我们来测试一下：
patches_embedded = PatchEmbedding()(x)
TransformerEncoderBlock()(patches_embedded).shape
output
Transformer
在ViT中只使用了原始Transformer中的Encoder部分。Encoder一共包含L个block，我们使用参数depth来指定，代码如下：
class TransformerEncoder(nn.Sequential):
def __init__(self, depth: int = 12, **kwargs):
super().__init__(*[TransformerEncoderBlock(**kwargs) for _ in range(depth)])
很简单不是嘛？
MLP Head
class ClassificationHead(nn.Sequential):
def __init__(self, emb_size: int = 768, n_classes: int = 1000):
super().__init__(
Reduce('b n e -> b e', reduction='mean'),
nn.LayerNorm(emb_size),
nn.Linear(emb_size, n_classes))
我们将之前定义好的PatchEMbedding, TransformerEncoder,ClassificationHead整合起来，搭建出最终的ViT代码模型如下：
class ViT(nn.Sequential):
def __init__(self,
in_channels: int = 3,
patch_size: int = 16,
emb_size: int = 768,
img_size: int = 224,
depth: int = 12,
n_classes: int = 1000,
**kwargs):
super().__init__(
PatchEmbedding(in_channels, patch_size, emb_size, img_size),
TransformerEncoder(depth, emb_size=emb_size, **kwargs),
ClassificationHead(emb_size, n_classes)
)
我们可以使用torchsummary函数来计算参数量，输出如下：
参数量
与其他ViT实现代码相比，这个参数量是差不多的。
原文的代码仓库在https://github.com/FrancescoSaverioZuppichini/ViT。
参考
- https://towardsdatascience.com/implementing-visualttransformer-in-pytorch-184f9f16f632
- https://arxiv.org/abs/2010.11929