最全的卷积神经网络实现手写数字识别,由于版本的更新,好多函数被替代或简化,本文是根据最新版本的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)))