共计 2422 个字符，预计需要花费 7 分钟才能阅读完成。
测试结果
最后两行分别为预测类别与真实类别。
数据预览
这里的数据使用的是 mnist 数据集，大家可以将代码中的 DOWNLOAD_MNIST 值修改为 True 进行自动下载。
代码
import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision # 数据库模块
import matplotlib.pyplot as plt
# 训练整批数据多少次，这里为了节约时间，只训练一次
EPOCH=1
#每次批处理 50 个数据
BATCH_SIZE=50
#学习效率
LR=0.001
# 如果已经下载好了 mnist 数据就写上 False
DOWNLOAD_MNIST = False
# 训练的数据集:Mnist 手写数字
train_data=torchvision.datasets.MNIST(
#保存或提取数据集的位置
root=’./mnist/’,
#该数据是训练数据
train=True,
#转换 PIL.Image or numpy.ndarray 成 torch.FloatTensor (C x H x W), 训练的时候 normalize 成 [0.0, 1.0] 区间
transform=torchvision.transforms.ToTensor(),
#没下载就下载, 下载了就不用再下了
download=DOWNLOAD_MNIST,
)
# 绘制一下数据集
#黑色的地方的值都是 0, 白色的地方值大于 0.
print(train_data.train_data.size()) # (60000, 28, 28)
print(train_data.train_labels.size()) # (60000)
plt.imshow(train_data.train_data[2].numpy(), cmap=’gray’)
plt.title(‘%i’ % train_data.train_labels[2])
plt.show()
# 测试数据
test_data=torchvision.datasets.MNIST(root=’./mnist/’,train=False)
# 批训练 50samples,1 channel,28×28 (50, 1, 28, 28)
train_loader=Data.DataLoader(dataset=train_data,batch_size=BATCH_SIZE,shuffle=True)
# 这里只测试了前 2000 个
#特征
test_x=torch.unsqueeze(test_data.test_data,dim=1).type(torch.FloatTensor)[:2000]/255.
#标签
test_y=test_data.test_labels[:2000]
# 构建 CNN 模型
class CNN(nn.Module):
def __init__(self):
super(CNN,self).__init__()
#input shape(1,28,28)
self.conv1=nn.Sequential(
#卷积
nn.Conv2d(
in_channels=1,
out_channels=16,
#filter size
kernel_size=5,
#filter movement/step
stride=1,
#如果想要 con2d 出来的图片长宽没有变化,
#padding=(kernel_size-1)/ 2 当 stride=1
padding=2,
),
#output shape(16,28,28)
#激励函数
nn.ReLU(),
#池化
# 在 2 ×2 空间里向下采样,output shape(16,14,14)
nn.MaxPool2d(kernel_size=2),
)
#input shape(16,14,14)
self.conv2=nn.Sequential(
nn.Conv2d(16,32,5,1,2),
#output shape(32,14,14)
#激励函数
nn.ReLU(),
#output shape(32,7,7)
nn.MaxPool2d(2),
)
#全连接层——进行分类。这里将其分成了 10 类
self.out=nn.Linear(32*7*7,10)
def forward(self,x):
x=self.conv1(x)
x=self.conv2(x)
#展平多维的卷积图成 (batch_size,32*7*7)
x=x.view(x.size(0),-1)
output=self.out(x)
return output
cnn=CNN()
print(cnn)
# 训练
#优化器
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):
#将数据输入 nn 并且得到 output
output=cnn(b_x)
#计算 output 与真实值之间的误差
loss=loss_func(output,b_y)
#清空上一步残余更新参数值
optimizer.zero_grad()
#误差反向传播，让参数进行更新
loss.backward()
#将更新后的参数值施加到 nn 的 parameters 上
optimizer.step()
# 测试: 选取 10 个数据
test_output=cnn(test_x[:10])
pred_y=torch.max(test_output,1)[1].data.numpy().squeeze()
print(pred_y, ‘prediction number’)
print(test_y[:10].numpy(), ‘real number’)
# if __name__==’__main__’:
# print(“hello word”)
