Tutorial de construcción de un modelo Pytorch desde cero (3) Construcción de una red de transformadores

Prefacio Este artículo presenta el proceso básico de Transformer, dos implementaciones de block, varias implementaciones de Position Embebdding, la implementación de Encoder, los dos métodos de clasificación finales y la introducción del formato de datos más importante.

Este artículo es de la serie de resúmenes técnicos de la guía técnica CV de cuenta pública

Bienvenido a la guía técnica de CV de cuenta pública , que se centra en el resumen técnico de la visión artificial, el seguimiento de última tecnología, la interpretación de documentos clásicos y la información de contratación de CV.

Antes de hablar de cómo construirlo, repasemos la estructura de Transformer en visión artificial. Aquí hay un ejemplo del ViT más típico.

Como se muestra en la figura, para una imagen, primero divídala en parches NxN, alise los parches, luego mapéelos en tokens a través de una capa completamente conectada y agregue incrustación de posición a cada token, que inicializará aleatoriamente tokens, después de concatenar para los tokens generados por la imagen, y luego pasando por el módulo codificador del transformador, después de pasar por múltiples capas de codificadores, los últimos tokens (es decir, tokens inicializados aleatoriamente) se eliminan, y luego la capa completamente conectada se usa como un red de clasificación para la clasificación.

A continuación, presentaremos paso a paso cómo construir un modelo de Transformador de acuerdo con este proceso. ,

Cuadra

En la actualidad, hay dos formas de lograr el bloqueo, una es la segmentación directa y la otra es a través de la convolución con el núcleo de convolución y el tamaño de zancada del tamaño del parche.

segmentación directa

La segmentación directa consiste en dividir directamente la imagen en varios bloques. La biblioteca einops debe usarse en la implementación del código. La operación completa es ajustar la forma de (B, C, H, W) a (B, (H/P *W/P), P*P*C) .

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

self.to_patch_embedding = nn.Sequential(
           Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
           nn.Linear(patch_dim, dim),
      )
复制代码

Aquí hay una breve introducción a Reorganizar.

Reorganizar se usa para volver a transformar y ordenar las dimensiones de los tensores, y se puede usar para reemplazar operaciones como reformar, ver, transponer y permutar en pytorch. para dar algunos ejemplos

#假设images的shape为[32,200,400,3]
#实现view和reshape的功能
Rearrange(images,'b h w c -> (b h) w c')#shape变为（32*200, 400, 3）
#实现permute的功能
Rearrange(images, 'b h w c -> b c h w')#shape变为（32, 3, 200, 400）
#实现这几个都很难实现的功能
Rearrange(images, 'b h w c -> (b c w) h')#shape变为（32*3*400, 200）
复制代码

从这几个例子看可以看出，Rearrange非常简单好用，这里的b, c, h, w都可以理解为表示符号，用来表示操作变化。通过这几个例子似乎也能理解下面这行代码是如何将图像分割的。

Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width)
复制代码

这里需要解释的是，一个括号内的两个变量相乘表示的是该维度的长度，因此不要把"h"和"w"理解成图像的宽和高。这里实际上h = H/p1, w = W/p2，代表的是高度上有几块，宽度上有几块。h和w都不需要赋值，代码会自动根据这个表达式计算，b和c也会自动对应到输入数据的B和C。

后面的"b (h w) (p1 p2 c)"表示了图像分块后的shape: （B，(H/P *W/P)，P*P*C）

这种方式在分块后还需要通过一层全连接层将分块的向量映射为tokens。

在ViT中使用的就是这种直接分块方式。

卷积分割

卷积分割比较容易理解，使用卷积核和步长都为patch大小的卷积对图像卷积一次就可以了。

self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)

x = self.proj(x).flatten(2).transpose(1, 2)  # B Ph*Pw C
复制代码

在swin transformer中即使用的是这种卷积分块方式。在swin transformer中卷积后没有再加全连接层。

Position Embedding

Position Embedding可以分为absolute position embedding和relative position embedding。

在学习最初的transformer时，可能会注意到用的是正余弦编码的方式，但这只适用于语音、文字等1维数据，图像是高度结构化的数据，用正余弦不合适。

在ViT和swin transformer中都是直接随机初始化一组与tokens同shape的可学习参数，与tokens相加，即完成了absolute position embedding。

在ViT中实现方式：

self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
x += self.pos_embedding[:, :(n + 1)]
#之所以是n+1，是因为ViT中选择随机初始化一个class token，与分块得到的tokens拼接。所以patches的数量为num_patches+1。
复制代码

在swin transformer中的实现方式：

from timm.models.layers import trunc_normal_
self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
trunc_normal_(self.absolute_pos_embed, std=.02)
复制代码

在TimeSformer中的实现方式：

self.pos_emb = torch.nn.Embedding(num_positions + 1, dim)
复制代码

以上就是简单的使用方法，这种方法属于absolute position embedding。

还有更复杂一点的方法，以后有机会单独搞一篇文章来介绍。

感兴趣的读者可以先去看看这篇论文《ICCV2021 | Vision Transformer中相对位置编码的反思与改进》。

Encoder

Encoder由Multi-head Self-attention和FeedForward组成。

Multi-head Self-attention

Multi-head Self-attention主要是先把tokens分成q、k、v，再计算q和k的点积，经过softmax后获得加权值，给v加权，再经过全连接层。

用公式表示如下：

所谓Multi-head是指把q、k、v再dim维度上分成head份，公式里的dk为每个head的维度。

具体代码如下：

class Attention(nn.Module):
   def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
       super().__init__()
       inner_dim = dim_head *  heads
       project_out = not (heads == 1 and dim_head == dim)

       self.heads = heads
       self.scale = dim_head ** -0.5
       self.attend = nn.Softmax(dim = -1)
       self.dropout = nn.Dropout(dropout)

       self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
       self.to_out = nn.Sequential(
           nn.Linear(inner_dim, dim),
           nn.Dropout(dropout)
      ) if project_out else nn.Identity()

   def forward(self, x):
       qkv = self.to_qkv(x).chunk(3, dim = -1)
       q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
       dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
       attn = self.attend(dots)
       attn = self.dropout(attn)

       out = torch.matmul(attn, v)
       out = rearrange(out, 'b h n d -> b n (h d)')
       return self.to_out(out)
复制代码

这里没有太多可以解释的地方，介绍一下q、k、v的来源，由于这是self-attention，因此q=k=v(即tokens)，若是普通attention，则k= v，而q是其它的东西，例如可以是另一个尺度的tokens，或视频领域中的其它帧的tokens。

FeedForward

这里不用多介绍。

class FeedForward(nn.Module):
   def __init__(self, dim, hidden_dim, dropout = 0.):
       super().__init__()
       self.net = nn.Sequential(
           nn.Linear(dim, hidden_dim),
           nn.GELU(),
           nn.Dropout(dropout),
           nn.Linear(hidden_dim, dim),
           nn.Dropout(dropout)
      )
   def forward(self, x):
       return self.net(x)
复制代码

把上面两者组合起来就是Encoder了。

class Transformer(nn.Module):
   def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
       super().__init__()
       self.layers = nn.ModuleList([])
       for _ in range(depth):
           self.layers.append(nn.ModuleList([
               PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
               PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
          ]))
   def forward(self, x):
       for attn, ff in self.layers:
           x = attn(x) + x
           x = ff(x) + x
       return x
复制代码

depth指的是Encoder的数量。PreNorm指的是层归一化。

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)
复制代码

分类方法

数据通过Encoder后获得最后的预测向量的方法有两种典型。在ViT中是随机初始化一个cls_token，concate到分块后的token后，经过Encoder后取出cls_token，最后将cls_token通过全连接层映射到最后的预测维度。

#生成cls_token部分
from einops import repeat
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))

cls_tokens = repeat(self.cls_token, '1 n d -> b n d', b = b)
x = torch.cat((cls_tokens, x), dim=1)
################################
#分类部分
self.mlp_head = nn.Sequential(
           nn.LayerNorm(dim),
           nn.Linear(dim, num_classes)
      )
x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

x = self.to_latent(x)
return self.mlp_head(x)
复制代码

在swin transformer中，没有选择cls_token。而是直接在经过Encoder后将所有数据取了个平均池化，再通过全连接层。

self.avgpool = nn.AdaptiveAvgPool1d(1)
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()

x = self.avgpool(x.transpose(1, 2))  # B C 1
x = torch.flatten(x, 1)
x = self.head(x)
复制代码

组合以上这些就成了一个完整的模型

class ViT(nn.Module):
   def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
       super().__init__()
       image_height, image_width = pair(image_size)
       patch_height, patch_width = pair(patch_size)

       num_patches = (image_height // patch_height) * (image_width // patch_width)
       patch_dim = channels * patch_height * patch_width
       assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

       self.to_patch_embedding = nn.Sequential(
           Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
           nn.Linear(patch_dim, dim),
      )

       self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
       self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
       self.dropout = nn.Dropout(emb_dropout)
       self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

       self.pool = pool
       self.to_latent = nn.Identity()
       self.mlp_head = nn.Sequential(
           nn.LayerNorm(dim),
           nn.Linear(dim, num_classes)
      )

   def forward(self, img):
       x = self.to_patch_embedding(img)
       b, n, _ = x.shape

       cls_tokens = repeat(self.cls_token, '1 n d -> b n d', b = b)
       x = torch.cat((cls_tokens, x), dim=1)
       x += self.pos_embedding[:, :(n + 1)]
       x = self.dropout(x)
       x = self.transformer(x)
       x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

       x = self.to_latent(x)
       return self.mlp_head(x)
复制代码

数据的变换

以上的代码都是比较简单的，整体上最麻烦的地方在于理解数据的变换。

首先输入的数据为（B, C, H, W），在经过分块后，变成了（B, n, d）。

在CNN模型中，很好理解(H，W)就是feature map，C是指feature map的数量，那这里的n，d哪个是通道，哪个是图像特征？

回顾一下分块的部分

Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width)
复制代码

根据这个可以知道n为分块的数量，d为每一块的内容。因此，这里的n相当于CNN模型中的C，而d相当于features。

一般情况下，在Encoder中，我们都是以（B, n, d）的形式。

在swin transformer中这种以卷积的形式分块，获得的形式为（B, C, L），然后做了一个transpose得到（B, L, C）,这与ViT通过直接分块方式获得的形式实际上完全一样，在Swin transformer中的L即为ViT中的n，而C为ViT中的d。

因此，要注意的是在Multi-head self-attention中，数据的形式是（Batchsize, Channel, Features），分成多个head的是Features。

前面提到，在ViT中会concate一个随机生成的cls_token，该cls_token的维度即为（B, 1, d）。可以理解为通道数多了个1。

以上就是Transformer的模型搭建细节了，整体上比较简单，大家看完这篇文章后可以找几篇Transformer的代码来理解理解。如ViT, swin transformer, TimeSformer等。

ViT：https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
swin: https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
TimeSformer:https://github.com/lucidrains/TimeSformer-pytorch/blob/main/timesformer_pytorch/timesformer_pytorch.py
复制代码

下一篇我们将介绍如何写train函数，以及包括设置优化方式，设置学习率，不同层设置不同学习率，解析参数等。

欢迎关注公众号CV技术指南，专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

CV技术指南创建了一个交流氛围很不错的群，除了太偏僻的问题，几乎有问必答。关注公众号添加编辑的微信号可邀请加交流群。