在这篇博客中，我们将详细介绍如何使用 PyTorch 进行 CIFAR-10 数据集的加载、训练和测试。我们将从数据预处理开始，介绍如何定义自定义数据集类和卷积神经网络，并最终展示如何训练模型和评估模型的性能。
1. 数据加载和预处理
首先，我们需要加载 CIFAR-10 数据集。CIFAR-10 是一个包含 60,000 张 32x32 彩色图像的数据集，共分为 10 类，每类 6,000 张图像。我们可以使用 torchvision.datasets 方便地下载和加载数据。
import torchvision.datasets as datasets
import numpy as np
trainset = datasets.CIFAR10(root='./data', train=True, download=True)
testset = datasets.CIFAR10(root='./data', train=False, download=True)
train_data = trainset.data
train_label = np.array(trainset.targets)
test_data = testset.data
test_label = np.array(testset.targets)
np.save('train_data.npy', train_data)
np.save('train_label.npy', train_label)
np.save('test_data.npy', test_data)
np.save('test_label.npy', test_label)
在这段代码中，我们下载了 CIFAR-10 数据集，并将训练和测试数据分别保存到 .npy 文件中，以便后续加载和处理。
2. 定义自定义数据集类
为了将数据加载到 PyTorch 中进行训练，我们需要定义一个自定义数据集类，继承自 torch.utils.data.Dataset。
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset
class CIFAR10Dataset(Dataset):
def __init__(self, transform, data, label):
super(CIFAR10Dataset, self).__init__()
self.transform = transform
self.images = data
self.labels = label
def __getitem__(self, idx):
img = self.images[idx]
img = self.transform(img)
label = torch.tensor(self.labels[idx], dtype=torch.long)
return img, label
def __len__(self):
return len(self.images)
这个自定义数据集类通过实现 __getitem__ 和 __len__ 方法来返回图像和标签，并进行必要的转换。
3. 定义卷积神经网络模型
接下来，我们定义一个简单的卷积神经网络（CNN）模型。该模型包含两个卷积层、两个池化层和三个全连接层。
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x
4. 训练模型
定义好数据集和模型后，我们可以开始训练模型。在训练之前，我们需要定义数据转换、加载数据、实例化模型、定义损失函数和优化器。
import torch.optim as optim
from torch.utils.data import DataLoader
def main():
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_data = np.load('train_data.npy')
train_label = np.load('train_label.npy')
test_data = np.load('test_data.npy')
test_label = np.load('test_label.npy')
trainset = CIFAR10Dataset(transform=transform, data=train_data, label=train_label)
trainloader = DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
testset = CIFAR10Dataset(transform=transform, data=test_data, label=test_label)
testloader = DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
net = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10):
running_loss = 0.0
for i, data in enumerate(trainloader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 2000 == 1999:
print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
running_loss = 0.0
print('Finished Training')
PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)
net = Net()
net.load_state_dict(torch.load(PATH))
5. 测试模型
训练完成后，我们可以测试模型的性能，计算每个类别的准确率
class_correct = list(0. for i in range(10))
class_total = list(0. for i in range(10))
classes = ['plane', 'car', 'bird', 'cat', 'deer',
'dog', 'frog', 'horse', 'ship', 'truck']
with torch.no_grad():
for data in testloader:
images, labels = data
outputs = net(images)
_, predicted = torch.max(outputs, 1)
c = (predicted == labels).squeeze()
for i in range(4):
label = labels[i]
class_correct[label] += c[i].item()
class_total[label] += 1
for i in range(10):
print('Accuracy of %5s : %2d %%' % (classes[i], 100 * class_correct[i] / class_total[i]))
if __name__ == '__main__':
main()
运行结果：
结论
这篇博客详细介绍了如何使用 PyTorch 加载 CIFAR-10 数据集，定义自定义数据集类，构建卷积神经网络模型，并进行训练和测试。通过以上步骤，你可以理解如何在实际项目中应用 PyTorch 进行图像分类任务。希望这篇博客对你有所帮助！
评论区