使用PyTorch框架构建一个基于ResNet50的10分类模型并进行训练,需要首先确保已经安装了PyTorch和必要的库(如torchvision,用于加载预训练的ResNet50模型)。以下是一个简单的步骤指导,包括模型构建、数据加载、训练循环和测试过程。
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F
假设的数据集是CIFAR-10(或其他类似的10分类数据集),需要定义合适的预处理步骤。对于ResNet50,由于CIFAR-10的图像较小(32x32),我们可能需要调整图像大小(虽然ResNet50是为224x224图像设计的)。
# 假设使用CIFAR-10
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
trainset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
testset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = DataLoader(testset, batch_size=64, shuffle=False)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
加载预训练的ResNet50,并修改最后的全连接层以匹配10个类别。
model = models.resnet50(pretrained=True)
# 修改最后的全连接层
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)
# 转到gpu
model.to('cuda')
# 将模型设置为训练模式
model.train()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 训练模型
num_epochs = 1
for epoch in range(num_epochs):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
# 梯度置零
optimizer.zero_grad()
# 前向传播
outputs = model(inputs.to('cuda'))# 转到gpu
loss = criterion(outputs, labels.to('cuda'))# 转到gpu
# 反向传播和优化
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 100 == 99: # 每100个mini-batches打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 100:.3f}')
running_loss = 0.0
print('Finished Training')
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = model(images.to('cuda')) # 转到gpu
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted.to('cpu') == labels).sum().item()
print(f'Accuracy of the network on the 10000 test images: {100 * correct / total} %')
以上步骤构建了一个基于ResNet50的10分类模型,并在CIFAR-10数据集上进行了训练和评估。注意,由于CIFAR-10的图像尺寸较小,我们进行了图像大小的调整以匹配ResNet50的输入要求。在实际应用中,可能需要根据具体的数据集和需求来调整这些步骤。