使用 Python 实现一个简单的图像分类器

在现代人工智能和计算机视觉领域，图像分类是一个非常基础且重要的任务。图像分类的目标是根据输入图像的内容将其分配到预定义的类别中。近年来，深度学习技术的发展使得图像分类的准确率得到了极大的提升。

本文将介绍如何使用 Python 和深度学习框架 PyTorch 来构建一个简单的图像分类器。我们将使用经典的 CIFAR-10 数据集进行训练和测试。整个过程包括数据加载、模型构建、训练、评估以及预测几个步骤。

环境准备

首先，确保你已经安装了以下库：

你可以通过以下命令安装这些依赖：

pip install torch torchvision matplotlib

步骤 1：导入必要的库

import torchimport torch.nn as nnimport torch.optim as optimimport torchvisionimport torchvision.transforms as transformsimport matplotlib.pyplot as plt

步骤 2：加载和预处理数据

CIFAR-10 是一个包含 60,000 张 32x32 彩色图像的数据集，分为 10 个类别（如飞机、汽车、鸟等）。我们将使用 torchvision 提供的 API 来加载数据，并对图像进行标准化预处理。

# 定义图像变换transform = transforms.Compose([    transforms.ToTensor(),    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])# 加载训练集和测试集trainset = torchvision.datasets.CIFAR10(root='./data', train=True,                                        download=True, transform=transform)trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,                                          shuffle=True, num_workers=2)testset = torchvision.datasets.CIFAR10(root='./data', train=False,                                       download=True, transform=transform)testloader = torch.utils.data.DataLoader(testset, batch_size=4,                                         shuffle=False, num_workers=2)classes = ('plane', 'car', 'bird', 'cat',           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

我们可以可视化一些训练图像来确认数据是否正确加载。

def imshow(img):    img = img / 2 + 0.5     # 反标准化    npimg = img.numpy()    plt.imshow(np.transpose(npimg, (1, 2, 0)))    plt.show()# 获取一批训练数据dataiter = iter(trainloader)images, labels = next(dataiter)# 显示图像imshow(torchvision.utils.make_grid(images))# 打印标签print(' '.join(f'{classes[labels[j]]}' for j in range(4)))

步骤 3：定义神经网络模型

我们将构建一个简单的卷积神经网络（CNN），包含两个卷积层和三个全连接层。

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(torch.relu(self.conv1(x)))        x = self.pool(torch.relu(self.conv2(x)))        x = x.view(-1, 16 * 5 * 5)        x = torch.relu(self.fc1(x))        x = torch.relu(self.fc2(x))        x = self.fc3(x)        return xnet = Net()

步骤 4：定义损失函数和优化器

我们使用交叉熵损失函数（Cross-Entropy Loss）和随机梯度下降（SGD）作为优化器。

criterion = nn.CrossEntropyLoss()optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

步骤 5：训练模型

接下来，我们开始训练模型。我们将训练 2 轮（epoch）以节省时间，但在实际应用中可以增加轮数以获得更好的性能。

for epoch in range(2):  # 多次循环遍历数据集    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:    # 每 2000 个小批量打印一次            print(f'Epoch {epoch + 1}, Batch {i + 1} loss: {running_loss / 2000:.3f}')            running_loss = 0.0print('Finished Training')

训练完成后，我们可以保存模型以便后续使用。

PATH = './cifar_net.pth'torch.save(net.state_dict(), PATH)

步骤 6：测试模型

现在我们来测试模型在测试集上的表现。

correct = 0total = 0with torch.no_grad():    for data in testloader:        images, labels = data        outputs = net(images)        _, predicted = torch.max(outputs.data, 1)        total += labels.size(0)        correct += (predicted == labels).sum().item()print(f'Accuracy of the network on the 10000 test images: {100 * correct / total:.2f}%')

步骤 7：查看每个类别的准确率

为了更细致地分析模型的表现，我们可以分别查看每个类别的准确率。

class_correct = list(0. for i in range(10))class_total = list(0. for i in range(10))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] += 1for i in range(10):    print(f'Accuracy of {classes[i]} : {100 * class_correct[i] / class_total[i]:.2f}%')

总结

本文介绍了如何使用 PyTorch 构建一个简单的图像分类器。我们从数据加载、模型构建、训练、评估到最终的预测都进行了详细说明，并提供了完整的代码示例。虽然这个模型相对简单，但它为我们理解图像分类的基本流程提供了一个很好的起点。

在实际项目中，通常会使用更复杂的网络结构（如 ResNet、VGG、EfficientNet 等）、数据增强技术、迁移学习等方法来进一步提高模型的性能。希望你能在此基础上继续探索深度学习的更多可能性！

参考文献：