本文最后更新于：14 天前

【深度学习】ViT入门（1）MNIST数据集练手

简介：视觉领域的Transformer —— ViT

什么是ViT

默子要和大家分享一款在计算机视觉领域非常火热的技术——ViT，也就是Vision Transformer的缩写。

虽然是几年前的图了，但这个图很明显，ViT强于传统的ResNet等一系列模型。

我们今天要学习的，就是 ViT

相信大家对Transformer这个名词并不陌生吧，它是一种革命性的深度学习模型，在自然语言处理领域取得了巨大的成功。

那么，什么是ViT呢？简单来说，ViT是一种基于Transformer架构的模型，它能够处理图像数据。

通常情况下，在处理图像时会使用卷积神经网络（Convolutional Neural Network，CNN），但是ViT却采用了和NLP领域中处理序列数据的Transformer相同的思路。

Transformer是一种基于自注意力机制（Self-Attention）的深度学习模型。

它不同于传统的卷积或者循环神经网络，而是通过全局的自注意力机制来学习序列之间的关系。

这种注意力机制允许模型在处理输入序列时能够对序列中不同位置的元素赋予不同的权重，从而捕捉到更丰富的上下文信息。

在ViT中，我们将图像视为一种特殊的序列数据，其中每个位置代表图像的一个小块区域，被称为补丁(patch)。这些补丁将被展平成向量形式，并作为输入序列传递给Transformer模型。

通过这种方式，ViT能够捕捉到图像中不同区域之间的上下文关系，从而实现对图像的理解和分析。

ViT训练过程

ViT模型的训练分为两个阶段：预训练和微调。在预训练阶段，模型会使用大规模的图像数据进行自监督学习，通过学习图像的不同补丁之间的关系来构建视觉表示。

在微调阶段，使用有标签的数据对模型进行进一步的训练，以适应具体的视觉任务，如图像分类、目标检测等。

优点

ViT在计算机视觉领域取得了惊人的成绩，甚至在一些任务上超越了传统的卷积神经网络。它的优势之一是能够处理全局的上下文信息，而不仅仅局限于局部区域。

此外，ViT还具有良好的可解释性，可以通过可视化自注意力权重来理解模型在图像中关注的区域。这对于理解模型的决策过程和解释结果非常有帮助。

此外，ViT还能够处理不同尺寸和分辨率的图像，而无需进行特定的调整或修改网络结构。这种灵活性使得ViT在处理多样化的图像数据时更加方便。

限制与挑战

当然，ViT也有一些挑战和限制。由于图像在序列化过程中丢失了空间信息，ViT可能无法充分捕捉到像素级别的细节。此外，ViT的计算成本较高，尤其是对于大型图像和复杂任务来说，需要更多的计算资源和时间。

不过，尽管存在这些挑战，ViT作为一种全新的视觉模型正在不断演进和发展。研究人员和工程师们正在努力改进ViT的性能，并将其应用于各种计算机视觉任务中。

在MNIST数据集上实现ViT

如何在MNIST数据集上实现ViT，默子将通过一个简单的例子来介绍。

提示：代码肯定跑不起来，但是可以作为参考。后面有放参考资料

首先，我们需要导入所需的库。我们将使用PyTorch来构建和训练我们的ViT模型，同时还需要导入一些辅助库来处理数据和可视化结果。

# 导入PyTorch相关库
import torch
# 导入PyTorch的神经网络模块
import torch.nn as nn
# 导入PyTorch的优化模块
import torch.optim as optim
# 导入PyTorch的数据加载模块
from torch.utils.data import DataLoader
# 导入PyTorch的视觉模块
from torchvision import datasets, transforms
# 导入Matplotlib绘图库
import matplotlib.pyplot as plt

接下来，我们定义一些超参数，例如批大小、学习率、迭代次数等。你可以根据需要进行调整。

1
2
3

batch_size = 64  # 每个批次的样本数量
learning_rate = 0.001  # 学习率
num_epochs = 10  # 训练的轮数（迭代次数）

现在，我们需要准备我们的数据集。MNIST数据集包含手写数字的灰度图像，每个图像都是28x28像素。我们将使用torchvision库来加载和预处理数据集。

transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为Tensor类型
    transforms.Normalize((0.5,), (0.5,))  # 对图像进行标准化处理
])

train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)  # 加载训练数据集
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)  # 加载测试数据集

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)  # 创建训练数据加载器
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)  # 创建测试数据加载器

现在我们来定义我们的ViT模型。在这个例子中，我们简化了ViT模型的结构，仅使用了一个Transformer编码器和一个全连接层作为输出。

class ViT(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes, num_heads, num_layers):
        super(ViT, self).__init__()
        self.embedding = nn.Linear(input_dim, hidden_dim)  # 输入层，将输入数据进行线性变换
        self.transformer_encoder = nn.TransformerEncoder(nn.TransformerEncoderLayer(hidden_dim, num_heads), num_layers)  # Transformer编码器层
        self.fc = nn.Linear(hidden_dim, num_classes)  # 全连接层，将隐藏层输出映射到类别数量上
        
    def forward(self, x):
        x = self.embedding(x)  # 嵌入层
        x = x.permute(1, 0, 2)  # 重新排列张量的维度，将batch_size和seq_len交换位置 (seq_len, batch_size, hidden_dim)
        output = self.transformer_encoder(x)  # Transformer编码器层的输出
        output = output.mean(dim=0)  # 沿着序列长度的维度进行平均
        output = self.fc(output)  # 全连接层的输出
        return output

接下来，我们实例化我们的ViT模型，并定义损失函数和优化器。

input_dim = 28 * 28  # 输入维度为图像的大小（28x28）
hidden_dim = 256  # 隐藏层维度
num_classes = 10  # 类别数量（0-9）
num_heads = 8  # Transformer中的注意力头数
num_layers = 4  # Transformer编码器层数

model = ViT(input_dim, hidden_dim, num_classes, num_heads, num_layers)  # 创建ViT模型实例
criterion = nn.CrossEntropyLoss()  # 定义损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # 定义优化器，将模型参数传入进行优化

现在我们可以开始训练我们的模型了。我们将使用训练集进行模型的训练，并在每个epoch结束后使用测试集评估模型的性能。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
model.to(device)  # 将模型移动到相应的设备上

for epoch in range(num_epochs):
    train_loss = 0.0
    train_correct = 0
    
    model.train()  # 设置模型为训练模式
    for images, labels in train_loader:
        images = images.view(-1, input_dim).to(device)  # 转换输入数据的维度并将其移动到设备上
        labels = labels.to(device)
        
        optimizer.zero_grad()  # 梯度置零
        
        outputs = model(images)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新模型参数
        
        train_loss += loss.item() * images.size(0)
        _, predicted = torch.max(outputs.data, 1)
        train_correct += (predicted == labels).sum().item()
    
    train_loss = train_loss / len(train_loader.dataset)
    train_accuracy = train_correct / len(train_loader.dataset)
    
    test_loss = 0.0
    test_correct = 0
    
    model.eval()  # 设置模型为评估模式
    with torch.no_grad():
        for images, labels in test_loader:
            images = images.view(-1, input_dim).to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            test_loss += loss.item() * images.size(0)
            _, predicted = torch.max(outputs.data, 1)
            test_correct += (predicted == labels).sum().item()
    
    test_loss = test_loss / len(test_loader.dataset)
    test_accuracy = test_correct / len(test_loader.dataset)
    
    print(f"Epoch {epoch+1}/{num_epochs}: Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.4f}, Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.4f}")

print("Training finished!")