pytorch+CNN经典模型→mnist识别

参考书籍：《深度学习框架pytorch快速开发与实战》

1、导入常用包

import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt

2、设置超参数

BATCH_SIZE = 50
EPOCH = 3
LR = 0.001

3、数据集下载及处理

（转化为pytorch处理的tensor格式）

train_data = torchvision.datasets.MNIST(
    root = './mnist',
    train = True,
    transform = torchvision.transforms.ToTensor(),
    download = True
)

test_data = torchvision.datasets.MNIST(
    root = './mnist',
    train = False,
    transform = torchvision.transforms.ToTensor(), 
    download = True
)

print(train_data.data.size()) 
print(test_data.data.size())

DataLoader可以把数据集分割为batch_size大小

train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

#test_x和test_y，维度上的一些处理
test_x = test_data.data.reshape(-1,1,28,28) #1是通道数，卷积核数量
test_x = torch.true_divide(test_x,255) #归一化，像素255
test_y = test_data.targets

4、模型构建

class CNN(nn.Module):
  def __init__(self):
    super(CNN,self).__init__(); #继承Module

    self.conv1 = nn.Sequential(
        nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2), #默认bias=True
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2) #pooling
    )

    self.conv2 = nn.Sequential(
        nn.Conv2d(16,32,5,1,2), #参见conv1中的参数顺序
        nn.ReLU(),
        nn.MaxPool2d(2)
    )

    self.out = nn.Linear(32*7*7, 10) #input_feature,out_feature,bias=True or False

  def forward(self,x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.size(0),-1) #linear的输入输出都是一维
    output = self.out(x)
    return output

5、实例化模型

cnn = CNN()
optimizer = torch.optim.Adam(cnn.parameters(), lr=LR)
loss_function = nn.CrossEntropyLoss() #交叉熵

6、训练模型、评价指标

for epoch in range(EPOCH):
  for step,(x,y) in enumerate(train_loader):
    b_x = Variable(x)
    b_y = Variable(y)

    output = cnn(b_x)
    loss = loss_function(output,y)
    optimizer.zero_grad() #梯度清0

    loss.backward()
    optimizer.step() #下一次更新

    if step%1000 == 0:
      test_output = cnn(test_x)
      pred_y = torch.max(test_output,1)[1].data.squeeze()
      accuracy =torch.true_divide(sum(pred_y==test_y), test_y.size(0))
      print('Epoch:',epoch,'|step:',step,'|train loss:%4f' % loss.item(),
            'test accuracy:%.4f' % accuracy)

7、看一下预测结果

test_output = cnn(test_x[:20])
# print(test_output)
pred_y = torch.max(test_output,1)[1].data.squeeze() #最大元素对应的下标
print(pred_y[:20],'prediction number')
print(test_y[:20],'real number')