关于最新版本GPU下pytorch实现卷积神经网络实现手写数字识别的代码

最全的卷积神经网络实现手写数字识别，由于版本的更新，好多函数被替代或简化，本文是根据最新版本的GPU11.0的pytorch撰写的卷积神经网络实现手写数字识别。欢迎大家借鉴。

 

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import torchvision
from torch.autograd import Variable
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

####定义卷积神经网络

#nn.Module,是专门为神经网络设计的模块接口
class CNN(nn.Module): #(相当于定义一个类)
  def __init__(self):
    #调用父类初始化函数(父类就是nn.Module)
    super(CNN, self).__init__()
#一个有序的容器，神经网络模块将按照传入的顺序依次添加到计算图中执行，
    self.conv1 = nn.Sequential(
#二卷积层，输入通道数是1，输出通道数是16，即16个卷积核
      nn.Conv2d(1, 6, 5, 1, 2), #1是一个通道即黑白图片，彩色图像是3，6是输出通道数，按自己需求决定，没有固定值
#BatchNorm2d是卷积网络中防止梯度消失或爆炸的函数，参数是卷积的输出通道数   (28*28-5+2*2)/1+1=784=28*28
      nn.BatchNorm2d(6),
      nn.ReLU(),
# 卷积后的图像大小：28*28*16,16为深度
      nn.MaxPool2d(2, 2)) 
# 池化后的图像大小：14*14*16

    self.conv2 = nn.Sequential(
      nn.Conv2d(6, 16, 5),
      nn.BatchNorm2d(16),
      nn.ReLU(),
#卷积后的图像大小：14*14*32,32为深度
      nn.MaxPool2d(2, 2))
#池化后的图像大小：7*7*32
    self.fc1 = nn.Sequential(
      nn.Linear(16 * 5 * 5, 120),
      nn.BatchNorm1d(120),
      nn.ReLU())

    self.fc2 = nn.Sequential(
      nn.Linear(120, 84),   ##这里需要计算，其他不需要， W2=（W1-F+2P）/S+1
      nn.BatchNorm1d(84),
      nn.ReLU(),
      nn.Linear(84, 10))
    	# 最后的结果一定要变为 10，因为数字的选项是 0 ~ 9

  def forward(self, x):
    x = self.conv1(x)
    x = self.conv2(x)
    x = x.view(x.size()[0], -1)  #为了适应全连接层的要求,将多维度的Tensor展平成一维，一行n列
    x = self.fc1(x)
    x = self.fc2(x)
    return x

####训练网络

batch_size = 64 #批处理的图片个数
learning_rate = 1e-2  #学习率，梯度下降算法的参数
num_epochs = 20   #全部训练集使用的次数
#下载数据后需要将格式转换成Tensor，且进行标准化
data_tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])

####下载数据集，位置放在本文件的父文件夹下的data文件夹里面

train_dataset = datasets.MNIST(root='./data', train=True, transform=data_tf, download=True)
test_dataset = datasets.MNIST(root='./data', train=False, transform=data_tf)
#DataLoader的作用是每个epoch的时候，对数据进行重新打乱，即重新分组,随机抽出一部分作为训练，每次的训练集都不一样
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
#对于测试集来说不需要进行从新分组
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

##看看train_loader中的随机抽取图片
for data,target in train_loader:
  for i in range(4):
    plt.figure()
#    print(target[i])
    plt.imshow(train_loader.dataset.data[i].numpy())

#如果cuda存在话，divice=‘cuda:0，否则是'cpu'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

net = CNN().to(device)   #放到GPU上运行，net就是类的实例化
# 损失函数使用交叉熵
criterion = nn.CrossEntropyLoss()
# 优化函数使用 Adam 自适应优化算法
optimizer = optim.Adam(
  net.parameters(),
  lr=learning_rate,
)

epoch = 1
net.train() #形成训练期间形成的网络，必加这句话
if __name__ == '__main__':
  for epoch in range(num_epochs):
    sum_loss = 0.0
    for i, data in enumerate(train_loader): # enumerate：提取序号和元素,i是序号，data是元素
      inputs, labels = data
      inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda()
# 根据pytorch中的backward()函数的计算，当网络参量进行反馈时，梯度是被积累的而不是被替换掉；
# 但是在每一个batch时毫无疑问并不需要将两个batch的梯度混合起来累积，因此这里就需要每个batch设置一遍zero_grad 了
# 将梯度初始化为零
      optimizer.zero_grad() #在反向传播之前要将梯度归零
##net(inputs)，即model(images)即等价于module.forward(images)  也就是将数据放到网络中去。
      outputs = net(inputs) #将数据传入网络进行前向运算，相当于调用forward函数，net.forward()，建立计算图
      loss = criterion(outputs, labels) #得到损失函数 ，相当于求出他两个的距离差距。
      loss.backward() #反向传播
      optimizer.step() #通过梯度做一步参数更新  主要对学习率自动的进行调整

      # print(loss)
      sum_loss += loss.item()
      if i % 100 == 99:      #每100次显示一下损失情况
        print('[%d,%d] loss:%.03f' %
           (epoch + 1, i + 1, sum_loss / 100))
        sum_loss = 0.0

#测试模型

  net.eval() #将模型变换为测试模式,为了固定BN和dropout层，使得偏置参数不随着发生变化，预测之前都要加上这句。
  correct = 0
  total = 0
  for data_test in test_loader:
    images, labels = data_test
    images, labels = Variable(images).cuda(), Variable(labels).cuda()
    output_test = net(images)
#函数返回两个tensor，第一个是每行的最大值，第二个是每行最大值的索引，因为softmax函数最大值是1
    _, predicted = torch.max(output_test, 1)  #找出最大的概率，意思就是这张图片倾向于这个类
    total += labels.size(0)
    correct += (predicted == labels).sum()  #预测结果是否等于目标
  print("correct1: ", correct)
  print("Test acc: {0}".format(correct.item() /
                 len(test_dataset)))