注意看代码注释！！！解析都在注释里，不要怕麻烦，自己复制代码，运行一遍，然后结合 运行结果和代码注释就能够看懂代码
1. 代码：
import os
# 引入一些包
import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt
# 定义一些参数
EPOCH = 1 # 训练数据的次数，我们这里假定训练一次
BATCH_SIZE = 50 # 每次训练的数据量，这个会产生每一次训练分多少次进行，或者多少批进行
LR = 0.001 # 学习率
DOWNLOAD_MNIST = False
# 下载并且加载数据集
if not(os.path.exists('./mnist/')) or not os.listdir('./mnist/'):
# not mnist dir or mnist is empyt dir
DOWNLOAD_MNIST = True
train_data = torchvision.datasets.MNIST(
root='./mnist/',
train=True, # 表示训练数据
transform=torchvision.transforms.ToTensor(), # 将数据转换成tensor
# torch.FloatTensor of shape (C x H x W) and normalize in the range [0.0, 1.0]
download=DOWNLOAD_MNIST,
)
# 画出一个例子，更直观
print(train_data.train_data.size()) # (60000, 28, 28)
print(train_data.train_labels.size()) # (60000)
plt.imshow(train_data.train_data[0].numpy(), cmap='gray')
plt.title('%i' % train_data.train_labels[0])
plt.show()
# 加载我们下载好的数据, 每一批的数据形状是 (50, 1, 28, 28)
train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True) # 加载数据
# 选择2000条数据来训练
test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000]/255. # shape from (2000, 28, 28) to (2000, 1, 28, 28), value in range(0,1)
test_y = test_data.test_labels[:2000]
# 构建自己的CNN网络
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Sequential( # 输入图形的形式 (1, 28, 28) 定义第一个卷积层
nn.Conv2d(
in_channels=1, # 输入的通道数，也就是高度
out_channels=16, # n_filters，16个过滤器 之后图形成了（16,28,28）
kernel_size=5, # 卷积核是5*5的
stride=1, # filter 过滤器的步长
padding=2, # if want same width and length of this image after Conv2d, padding=(kernel_size-1)/2 if stride=1
), # output shape (16, 28, 28)
nn.ReLU(), # activation 激活函数
nn.MaxPool2d(kernel_size=2), # 选择 2x2 area,进行池化层操作， 输出形状 (16, 14, 14)
)
self.conv2 = nn.Sequential( # 输入形状 (16, 14, 14)
nn.Conv2d(16, 32, 5, 1, 2), # 输出形状 (32, 14, 14)
nn.ReLU(), # 激活函数
nn.MaxPool2d(2), # 池化层之后的形状 (32, 7, 7)
)
self.out = nn.Linear(32 * 7 * 7, 10) # 全连接层, 输出10个数字，因为分类嘛，总共有10个类。
def forward(self, x):
x = self.conv1(x)
x = self.conv2(x)
x = x.view(x.size(0), -1) # 将数据由（32,7,7)这样的空间数据拉成一个列向量，也就是32*7*7
output = self.out(x)
return output, x # return x for visualization
cnn = CNN()
# 打印出来看一看
print(cnn) # net architecture
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR) # 在优化器中传入参数
loss_func = nn.CrossEntropyLoss() # 专门用来做分类的损失函数
# 开始训练
for epoch in range(EPOCH):
for step, (b_x, b_y) in enumerate(train_loader): # 数据总量/每批训练量=最终step的值
print('b_x: ',b_x)
output = cnn(b_x)[0] # cnn output
loss = loss_func(output, b_y) # cross entropy loss
optimizer.zero_grad() # clear gradients for this training step
loss.backward() # 神经网络反向传播
optimizer.step() # 更新梯度，或者更新参数
if step % 50 == 0:
test_output, last_layer = cnn(test_x)
pred_y = torch.max(test_output, 1)[1].data.numpy()
accuracy = float((pred_y == test_y.data.numpy()).astype(int).sum()) / float(test_y.size(0)) # 计算正确率
print('Epoch: ', epoch, '| train loss: %.4f' % loss.data.numpy(), '| test accuracy: %.2f' % accuracy)
# 使用10个测试数据进行测试
test_output, _ = cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.numpy()
print('pred_y_1: ',test_output)
print('pred_y_2: ',torch.max(test_output,1))
print('pred_y_3: ',torch.max(test_output,1)[1])
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')
2. 运行结果：
