PyTorch 代码了解 VIT Vision Transformer

2024-02-03 17:19:42

介绍

Vision Transformer (VIT) 是一种用于图像分类的深度学习模型。它是由谷歌大脑团队在 2020 年提出的，并在计算机视觉领域引起了广泛的关注。VIT 的主要思想是将图像转换为序列，然后使用 Transformer 模型进行处理。这种方法与传统的卷积神经网络 (CNN) 不同，CNN 直接在图像上进行操作。

VIT 的主要组件

VIT 由以下几个主要组件组成：

Transformer 编码器： Transformer 编码器是 VIT 的核心组件。它负责将图像转换为序列。Transformer 编码器由多个子层组成，包括自注意力机制、前馈神经网络和残差连接。
分类头： 分类头是一个全连接层，用于将 Transformer 编码器的输出分类为不同的类别。
位置嵌入： 位置嵌入是添加到图像序列中的特殊标记，以提供每个元素在序列中的位置信息。

代码示例

以下是一个使用 PyTorch 实现 VIT 的代码示例：

import torch
import torch.nn as nn
from einops import rearrange, repeat

class VIT(nn.Module):
    def __init__(self, num_classes, dim, depth, heads, mlp_dim, dropout = 0.1):
        super().__init__()
        self.num_classes = num_classes
        self.token_dim = dim
        self.vit = nn.ModuleList([nn.ModuleList([
            nn.Linear(dim, dim),
            nn.Linear(dim, dim),
        ]) for _ in range(depth)])
        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, 1 + num_classes, dim))
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        b, n, _ = x.shape
        cls_tokens = repeat(self.cls_token, '() -> b 1 d', b = b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embed[:, :(n + 1)]
        for blk in self.vit:
            x = self.attn(self.norm1(x)) + x
            x = self.mlp(self.norm2(x)) + x
        cls_token = x[:, 0]
        return self.head(cls_token)

    def norm1(self, x):
        return self.layer_norm1(x.transpose(1, 2)).transpose(1, 2)

    def norm2(self, x):
        return self.layer_norm2(x.transpose(1, 2)).transpose(1, 2)

    def mlp(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        return x

    def attn(self, x):
        x = self.qkv_proj(x)
        q, k, v = x.chunk(3, dim = -1)
        attn = q.matmul(k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim = -1)
        attn = self.dropout(attn)

        x = attn.matmul(v)
        return x.transpose(1, 2)

    def head(self, x):
        x = self.last_linear(x)
        return x