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

在现代机器学习和人工智能领域，图像分类是一项基础而重要的任务。本文将介绍如何使用 Python 和深度学习框架 PyTorch 来构建一个简单的图像分类器，该分类器能够识别手写数字（基于 MNIST 数据集）。

我们将从数据准备、模型构建、训练到评估整个流程进行详细讲解，并提供完整的代码示例。

1. 环境准备

首先，我们需要安装必要的库：

pip install torch torchvision matplotlib

2. 数据加载与预处理

我们使用经典的 MNIST 手写数字数据集，它包含 60,000 张训练图片和 10,000 张测试图片，每张图片是 28x28 的灰度图，表示数字 0 到 9。

import torchfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as plt# 数据预处理：将图像转换为张量并归一化transform = transforms.Compose([    transforms.ToTensor(),    transforms.Normalize((0.5,), (0.5,))])# 加载训练集和测试集train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)# 创建数据加载器batch_size = 64train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)# 可视化部分样本def show_images(loader):    dataiter = iter(loader)    images, labels = next(dataiter)    fig, axes = plt.subplots(2, 5, figsize=(10, 5))    for i, ax in enumerate(axes.flat):        ax.imshow(images[i].squeeze(), cmap='gray')        ax.set_title(f'Label: {labels[i]}')        ax.axis('off')    plt.show()show_images(train_loader)

运行上述代码后，我们可以看到一些训练样本的图像及其标签。

3. 构建神经网络模型

我们使用一个简单的卷积神经网络（CNN），结构如下：

import torch.nn as nnimport torch.nn.functional as Fclass SimpleCNN(nn.Module):    def __init__(self):        super(SimpleCNN, self).__init__()        # 第一层卷积层        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=1, padding=1)        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)        # 第二层卷积层        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1)        # 全连接层        self.fc1 = nn.Linear(32 * 7 * 7, 128)        self.fc2 = nn.Linear(128, 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, 32 * 7 * 7)        x = F.relu(self.fc1(x))        x = self.fc2(x)        return x

4. 训练模型

接下来我们定义损失函数和优化器，并开始训练模型。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = SimpleCNN().to(device)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=0.001)# 训练循环num_epochs = 5for epoch in range(num_epochs):    model.train()    running_loss = 0.0    for images, labels in train_loader:        images, labels = images.to(device), labels.to(device)        optimizer.zero_grad()        outputs = model(images)        loss = criterion(outputs, labels)        loss.backward()        optimizer.step()        running_loss += loss.item()    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')

5. 模型评估

训练完成后，我们在测试集上评估模型性能。

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

6. 可视化预测结果

为了进一步验证模型效果，我们可以随机选择几个测试图像并显示模型的预测结果。

def visualize_predictions():    model.eval()    dataiter = iter(test_loader)    images, labels = next(dataiter)    images, labels = images.to(device), labels.to(device)    outputs = model(images)    _, predicted = torch.max(outputs, 1)    fig, axes = plt.subplots(2, 5, figsize=(10, 5))    for i, ax in enumerate(axes.flat):        ax.imshow(images[i].cpu().squeeze(), cmap='gray')        ax.set_title(f'Pred: {predicted[i].item()}\nTrue: {labels[i].item()}')        ax.axis('off')    plt.tight_layout()    plt.show()visualize_predictions()

7. 总结

本文通过使用 PyTorch 实现了一个简单的 CNN 图像分类器，对 MNIST 手写数字进行了训练和评估。我们涵盖了以下关键技术点：

虽然这个模型相对简单，但它展示了图像分类的基本流程。读者可以在本项目基础上尝试更复杂的网络结构（如 ResNet、VGG）、使用更多数据增强技术或迁移学习来提升模型性能。

附录：完整代码汇总

以下是整合后的完整代码：

import torchfrom torchvision import datasets, transformsfrom torch.utils.data import DataLoaderimport matplotlib.pyplot as pltimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optim# 数据预处理transform = transforms.Compose([    transforms.ToTensor(),    transforms.Normalize((0.5,), (0.5,))])train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=True)test_dataset = datasets.MNIST(root='./data', train=False, transform=transform, download=True)batch_size = 64train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)# 显示图像def show_images(loader):    dataiter = iter(loader)    images, labels = next(dataiter)    fig, axes = plt.subplots(2, 5, figsize=(10, 5))    for i, ax in enumerate(axes.flat):        ax.imshow(images[i].squeeze(), cmap='gray')        ax.set_title(f'Label: {labels[i]}')        ax.axis('off')    plt.show()# 定义模型class SimpleCNN(nn.Module):    def __init__(self):        super(SimpleCNN, self).__init__()        self.conv1 = nn.Conv2d(1, 16, 3, padding=1)        self.pool = nn.MaxPool2d(2, 2)        self.conv2 = nn.Conv2d(16, 32, 3, padding=1)        self.fc1 = nn.Linear(32 * 7 * 7, 128)        self.fc2 = nn.Linear(128, 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, 32 * 7 * 7)        x = F.relu(self.fc1(x))        x = self.fc2(x)        return x# 初始化模型device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = SimpleCNN().to(device)criterion = nn.CrossEntropyLoss()optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练num_epochs = 5for epoch in range(num_epochs):    model.train()    running_loss = 0.0    for images, labels in train_loader:        images, labels = images.to(device), labels.to(device)        optimizer.zero_grad()        outputs = model(images)        loss = criterion(outputs, labels)        loss.backward()        optimizer.step()        running_loss += loss.item()    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {running_loss/len(train_loader):.4f}')# 评估model.eval()correct = 0total = 0with torch.no_grad():    for images, labels in test_loader:        images, labels = images.to(device), labels.to(device)        outputs = model(images)        _, predicted = torch.max(outputs.data, 1)        total += labels.size(0)        correct += (predicted == labels).sum().item()print(f'Accuracy on test set: {100 * correct / total:.2f}%')# 可视化预测def visualize_predictions():    model.eval()    dataiter = iter(test_loader)    images, labels = next(dataiter)    images, labels = images.to(device), labels.to(device)    outputs = model(images)    _, predicted = torch.max(outputs, 1)    fig, axes = plt.subplots(2, 5, figsize=(10, 5))    for i, ax in enumerate(axes.flat):        ax.imshow(images[i].cpu().squeeze(), cmap='gray')        ax.set_title(f'Pred: {predicted[i].item()}\nTrue: {labels[i].item()}')        ax.axis('off')    plt.tight_layout()    plt.show()visualize_predictions()

希望这篇文章能帮助你入门图像分类任务，并为进一步探索计算机视觉和深度学习打下坚实的基础！如果你有任何问题，欢迎留言讨论。