介绍
这段代码是使用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数据集上对手写数字进行分类，并在单独的测试数据集上评估其性能。
代码
python
# -*- 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 out
cnn = 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().size
print('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]