使用torchvision集成的efficientnet-v2-s模型,调用torchvision库中的Oxford IIIT Pet数据集,对模型进行训练。
若有修改要求,可以修改以下部分:
train_dataset = OxfordIIITPet(root='./data', split='trainval', download=True, transform=transform_train)
test_dataset = OxfordIIITPet(root='./data', split='test', download=True, transform=transform_test)
#常见数据集可以直接加载,若是自己的数据集就自己写个dataset/dataloader
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 37)
#37为数据集类别数,修改为自己对应的
scheduler = ReduceLROnPlateau(optimizer, 'max', patience=3, factor=0.1, verbose=True)
#学习率处可以自己调整,可玩性较高
训练截图:
其实十轮左右就稳定在90以上了,跑了三十轮,记得修改保存路径,我这里是用kaggle跑的。
代码如下:
import torch
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.models import efficientnet_v2_s
import torch.nn as nn
import torch.optim as optim
from torchvision.datasets import OxfordIIITPet
from torch.optim.lr_scheduler import ReduceLROnPlateau
from tqdm import tqdm
# 数据预处理 + 数据增强
transform_train = transforms.Compose([
transforms.Resize((256, 256)), # 增大图片预处理尺寸
transforms.RandomCrop((224, 224)), # 随机裁剪到模型输入尺寸
transforms.RandomHorizontalFlip(), # 随机水平翻转
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # 颜色抖动
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
transform_test = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 加载数据集
train_dataset = OxfordIIITPet(root='./data', split='trainval', download=True, transform=transform_train)
test_dataset = OxfordIIITPet(root='./data', split='test', download=True, transform=transform_test)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)
# 模型定义
model = efficientnet_v2_s(pretrained=True)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 37)
# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.001)
scheduler = ReduceLROnPlateau(optimizer, 'max', patience=3, factor=0.1, verbose=True)
# 训练模型
def train_model(num_epochs):
model.train()
best_accuracy = 0
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for inputs, labels in tqdm(train_loader, desc=f'Training Epoch {epoch + 1}'):
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
# 每个 epoch 后测试
accuracy = test_model()
scheduler.step(accuracy)
# 如果当前模型表现更好,保存模型
if accuracy > best_accuracy:
best_accuracy = accuracy
torch.save(model.state_dict(), '/kaggle/working/best_oxford_pets_efficientnetv2.pth')
print(f'New best model saved with accuracy: {best_accuracy:.2f}%')
def test_model():
model.eval()
correct = 0
total = 0
with torch.no_grad():
for inputs, labels in tqdm(test_loader, desc='Testing'):
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
# 调试输出
if total < 50: # 只打印前50个样本的信息
print(f'Predicted: {predicted[:10]}, Labels: {labels[:10]}')
accuracy = 100 * correct / total
print(f'Testing Accuracy: {accuracy:.2f}%')
return accuracy
train_model(30)