介绍
这段代码是使用PyTorch实现的卷积神经网络(CNN),用于在MNIST数据集上进行图像分类。让我一步步解释:
-
导入库:代码导入了必要的库,包括PyTorch(
torch
)、神经网络模块(torch.nn
)、函数模块(torch.nn.functional
)、图像数据集(torchvision
)以及数据处理(torch.utils.data
)和可视化(matplotlib.pyplot
)的工具。 -
设置超参数:定义了超参数,如批大小(
Batch_size
)、epoch数量(Epoch
)和学习率(Lr
)。 -
加载MNIST数据集:使用
torchvision.datasets.MNIST
加载MNIST数据集。该数据集包含了0到9的手写数字的灰度图像。transform=torchvision.transforms.ToTensor()
将PIL图像转换为PyTorch张量。 -
可视化样本数据:打印数据集的大小,并显示数据集中的第一张图像及其相应的标签。
-
准备测试数据:准备测试数据与训练数据类似。加载MNIST测试数据集,并选择前2000个图像进行测试。
-
创建数据加载器:使用
torch.utils.data.DataLoader
创建训练数据的数据加载器。它有助于在训练过程中对数据进行分批和混洗。 -
定义CNN架构:通过子类化
nn.Module
来定义CNN类。该架构包括两个卷积层(self.con1
、self.con2
),后面跟有ReLU激活函数和最大池化层。卷积层的输出被展平并馈入全连接层(self.out
),产生最终输出。 -
初始化CNN:创建CNN类的实例。
-
定义损失函数和优化器:使用交叉熵损失(
nn.CrossEntropyLoss
)作为损失函数,使用随机梯度下降(torch.optim.SGD
)作为优化器。 -
训练CNN:在指定的epoch数量循环内训练模型。在循环内,将训练数据通过模型,计算损失,进行梯度反向传播,并由优化器更新模型参数。
-
测试模型:每50次迭代训练时,对测试数据集进行评估。将测试预测与真实标签进行比较,计算准确率。
-
打印结果:训练结束后,打印模型预测及前10个测试样本的真实标签。
总的来说,这段代码训练了一个CNN模型,用于在MNIST数据集上对手写数字进行分类,并在单独的测试数据集上评估其性能。
代码
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torch.utils.data as Data
import matplotlib.pyplot as plt# define hyper parameters
Batch_size = 100
Epoch = 1
Lr = 0.5
#DOWNLOAD_MNIST = True # 若没有数据,用此生成数据# define train data and test data
train_data = torchvision.datasets.MNIST(root='./mnist',train=True,download=False,transform=torchvision.transforms.ToTensor()
)
print(train_data.data.size())
print(train_data.targets.size())
print(train_data.data[0])
# 画一个图片显示出来
plt.imshow(train_data.data[0].numpy(),cmap='gray')
plt.title('%i'%train_data.targets[0])
plt.show()
# print(train_data.data.shape) # torch.Size([60000, 28, 28])
# print(train_data.targets.size()) # torch.Size([60000])
# print(train_data.data[0].size()) # torch.Size([28, 28])
# plt.imshow(train_data.data[0].numpy(), cmap='gray')
# plt.show()
test_data = torchvision.datasets.MNIST(root='./mnist',train=False,# transform=torchvision.transforms.ToTensor()
)
test_x = torch.unsqueeze(test_data.data, dim=1).type(torch.FloatTensor)[:2000]
test_y = test_data.targets[:2000]
# print(test_x.shape) # torch.Size([2000, 1, 28, 28])
# print(test_y.shape) # torch.Size([2000])
train_loader = Data.DataLoader(dataset=train_data,shuffle=True,batch_size=Batch_size,
)# define network structure
class CNN(nn.Module):def __init__(self):super(CNN, self).__init__()self.con1 = nn.Sequential(nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),nn.ReLU(),nn.MaxPool2d(kernel_size=2))self.con2 = nn.Sequential(nn.Conv2d(16, 32, 5, 1, 2),nn.ReLU(),nn.MaxPool2d(2))self.out = nn.Linear(32 * 7 * 7, 10)def forward(self, x):x = self.con1(x) # (batch, 16, 14, 14)x = self.con2(x) # (batch, 32, 7, 7)x = x.view(x.size(0), -1)out = self.out(x) # (batch_size, 10)return outcnn = CNN()
# print(cnn)
optimizer = torch.optim.SGD(cnn.parameters(), lr=Lr)
loss_fun = nn.CrossEntropyLoss()for epoch in range(Epoch):for i, (x, y) in enumerate(train_loader):output = cnn(x)loss = loss_fun(output, y)optimizer.zero_grad()loss.backward()optimizer.step()if i % 50 == 0:test_output = torch.max(cnn(test_x), dim=1)[1]loss = loss_fun(cnn(test_x), test_y).item()accuracy = torch.sum(torch.eq(test_output, test_y)).item() / test_y.numpy().sizeprint('Epoch:', Epoch, '|loss:%.4f' % loss, '|accuracy:%.4f' % accuracy)print('real value', test_data.targets[: 10].numpy())
print('train value', torch.max(cnn(test_x)[: 10], dim=1)[1].numpy())
结果
real value [7 2 1 0 4 1 4 9 5 9]
train value [7 2 1 0 4 1 4 9 5 9]