博客

用Pytorch训练CNN(数据集MNIST,使用GPU的方法)

  • Post category:Python

以下是使用PyTorch训练CNN（数据集MNIST，使用GPU的方法）的完整攻略。

步骤一：导入必要的库

首先，我们需要导入必要的库，包括PyTorch、torchvision、numpy和matplotlib等。

import torch
import torchvision
import numpy as np
import matplotlib.pyplot as plt

步骤二：加载数据集

接下来，我们需要加载MNIST数据集。可以使用torchvision中的datasets模块来加载数据集。

train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=torchvision.transforms.ToTensor(), download=True)

步骤三：定义模型

我们使用一个简单的卷积神经网络来实现手写数字识别。定义模型的代码如下：

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=5, padding=2)
        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=5, padding=2)
        self.fc1 = torch.nn.Linear(7 * 7 * 64, 1024)
        self.fc2 = torch.nn.Linear(1024, 10)

    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = torch.nn.functional.max_pool2d(x, 2)
        x = torch.nn.functional.relu(self.conv2(x))
        x = torch.nn.functional.max_pool2d(x, 2)
        x = x.view(-1, 7 * 7 * 64)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.dropout(x, training=self.training)
        x = self.fc2(x)
        return torch.nn.functional.log_softmax(x, dim=1)

model = Net()

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

我们使用交叉熵损失函数和随机梯度下降优化器来训练模型。

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)

步骤五：将模型和数据移动到GPU上

我们需要将模型和数据移动到GPU上进行训练。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

步骤六：训练模型

接下来，我们使用训练集对模型进行训练。

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)

for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

步骤七：测试模型

最后，我们使用测试集对模型进行测试。

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=True)

with torch.no_grad():
    correct = 0
    total = 0
    for data, target in test_loader:
        data, target = data.to(device), target.to(device)
        output = model(data)
        _, predicted = torch.max(output.data, 1)
        total += target.size(0)
        correct += (predicted == target).sum().item()

    print('Accuracy of the network on the 10000 test images: %d %%' % (
        100 * correct / total))

上面的代码实现了使用PyTorch训练CNN（数据集MNIST，使用GPU的方法）。下面是两个示例：

示例一：显示数据集中的一张图片

image, label = train_dataset[0]
plt.imshow(image.squeeze().numpy(), cmap='gray')
plt.title('Label: %d' % label)
plt.show()

示例二：显示模型的预测结果

image, label = test_dataset[0]
image = image.to(device)
output = model(image.unsqueeze(0))
_, predicted = torch.max(output.data, 1)
plt.imshow(image.cpu().squeeze().numpy(), cmap='gray')
plt.title('Predicted: %d, Actual: %d' % (predicted.item(), label))
plt.show()

上面的代码分别显示了数据集中的一张图片和模型的预测结果。

Tags: