在深度学习领域，预训练模型已经成为提高模型性能和减少训练时间的重要手段。然而，有时候我们可能需要从头开始训练模型，以便更好地适应特定的数据集或者学习任务。本文将介绍如何从零开始训练 Vision Transformer (ViT) 模型进行图像分类任务。
Vision Transformer 简介
Vision Transformer (ViT) 是一种将自然语言处理中大放异彩的 Transformer 架构应用于计算机视觉任务的模型。ViT 通过将图像分割成多个小块（patches），然后将这些小块视为序列数据输入到 Transformer 模型中，从而实现对图像的分类和其他视觉任务。
环境准备
在开始之前，请确保您的环境中安装了以下库：
- PyTorch
- torchvision
- transformers
- PIL 或其他图像处理库
您可以通过 pip 安装这些库：
pip install torch torchvision transformers pillow
数据预处理
在训练模型之前，我们需要对图像数据进行预处理。这通常包括调整图像大小、随机翻转、旋转等操作，以增加数据的多样性并防止模型过拟合。
from torchvision import transforms
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.RandomRotation(20),
transforms.ToTensor(),
])
加载数据集
我们将使用 ImageFolder 来加载数据集，并将其分为训练集和测试集。
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
train_dataset = ImageFolder(root='path_to_train_dataset', transform=transform)
test_dataset = ImageFolder(root='path_to_test_dataset', transform=transform)
train_loader = DataLoader(dataset=train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=16, shuffle=False)
初始化 Vision Transformer 模型
我们将初始化一个 ViT 模型，不加载任何预训练权重。我们还需要修改模型的分类器层，以适应我们的二分类任务。
from transformers import ViTForImageClassification, ViTConfig
config = ViTConfig(
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
hidden_size=768,
num_classes=2 # 二分类
)
model = ViTForImageClassification(config)
训练模型
接下来，我们将定义损失函数、优化器和学习率调度器，并开始训练模型。
import torch
import torch.nn as nn
import torch.optim as optim
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=0.0001)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
# 训练循环
for epoch in range(20):
# 训练过程...
测试模型
在训练完成后，我们将在测试集上评估模型的性能。
model.eval()
# 测试过程...
保存模型
最后，我们将模型的权重保存到文件中，以便将来使用。
torch.save(model.state_dict(), 'path_to_save_model/transformer_model.pth')
结论
从头开始训练 Vision Transformer 模型可以让我们更好地控制模型的学习过程，并完全适应我们的特定任务。虽然这可能需要更多的时间和计算资源，但有时这是获得最佳性能的必要步骤。
通过本文的介绍，您应该能够理解并实现一个不使用预训练权重的 ViT 模型。开始训练您自己的模型，并探索深度学习在图像分类任务中的潜力吧！
transformer_train.py
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader
from transformers import ViTFeatureExtractor, ViTForImageClassification, ViTConfig
from torchvision.datasets import ImageFolder
import warnings
warnings.filterwarnings('ignore')
# 设置超参数
batch_size = 16
learning_rate = 0.0001
num_epochs = 20
num_classes = 2 # 二分类
# 数据预处理
transform = transforms.Compose([
transforms.Resize((224, 224)), # 调整图像大小以适应ViT模型
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.RandomRotation(20),
transforms.ToTensor(),
])
# 加载数据集
train_dataset = ImageFolder(root='E:/PycharmProject/LargeSoilDetection/datasets/train/', transform=transform)
test_dataset = ImageFolder(root='E:/PycharmProject/LargeSoilDetection/datasets/test/', transform=transform)
# 创建数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
# 初始化ViT模型配置
config = ViTConfig()
config.num_hidden_layers = 12 # 您可以根据需要调整模型大小
config.num_attention_heads = 12
config.intermediate_size = 3072
config.hidden_size = 768
config.num_classes = num_classes # 设置正确的类别数
# 初始化模型
model = ViTForImageClassification(config)
# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters(), lr=learning_rate)
# 学习率衰减
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
# 训练模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)
model.to(device)
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for images, labels in train_loader:
images = images.to(device)
labels = labels.to(device)
# 前向传播
outputs = model(images)
loss = criterion(outputs.logits, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
running_loss += loss.item()
epoch_loss = running_loss / len(train_loader)
print(f'Epoch [{
epoch+1}/{
num_epochs}], Loss: {
epoch_loss:.4f}')
# 测试模型
model.eval()
with torch.no_grad():
correct = 0
total = 0
for images, labels in test_loader:
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.logits, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
accuracy = 100 * correct / total
print(f'Accuracy of the model on the test images: {
accuracy:.2f}%')
# 保存模型
torch.save(model.state_dict(), './model/transformer_model2.pth')