CNN_for_MNIST，深度学习界的Hello World

github:https://github.com/Nand-Lu/cnn_for_mnist

导入模块

import torch 
import torch.nn as nn
import torch.utils.data as Data
import torchvision
import matplotlib.pyplot as plt
import os
import numpy as np

准备数据集

mnist_dowmload =False  

if not os.path.exists('./mnist/'): #判断路径下是否已经存在数据集
    mnist_dowmload = True #没有存在，改变布尔值为真
train_data = torchvision.datasets.MNIST( #MNITS数据集的加载函数
                                root= './mnist/',
                                train= True,#False为测试集
                                transform=torchvision.transforms.ToTensor(),
                                download=mnist_dowmload)

数据加载

train_loader =Data.DataLoader(dataset= train_data, batch_size=20, shuffle=True ) #训练集加载，batch_size根据自己显卡内存自己设置
test_data = torchvision.datasets.MNIST(root= './mnist/',train=False)#测试集
test_x = torch.unsqueeze(test_data.data ,dim=1).type(torch.FloatTensor)[:]
test_y = test_data.targets[:]

cnn网络

class cnn(nn.Module):
    
    def __init__(self):
        super(cnn, self).__init__()
        self.conv1 = nn.Sequential(#各种层的“容器”
        nn.Conv2d(in_channels= 1, out_channels= 16,kernel_size= 5,stride= 1,padding= 2),#卷积的参数设置
        nn.ReLU(),#激活函数
        nn.MaxPool2d(kernel_size=2)#最大池化
        )
        self.conv2 = nn.Sequential(
        nn.Conv2d(in_channels=16, out_channels= 32, kernel_size=5, stride= 1, padding= 2),
        nn.ReLU(),
        nn.MaxPool2d(kernel_size=2)
                             )
        self.out = nn.Linear(in_features=32*7*7,out_features=10) #全连接层输出
     
    def forward(self, x):#前向传播
        
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output, x
cnn = cnn()

优化器及损失函数

optimizer = torch.optim.Adam(cnn.parameters(),lr=0.01)#Adam优化器，学习率lr是超参数，可以自己调节
loss_fun = nn.CrossEntropyLoss()#损失函数

训练

Epoch = 1#超参数，自己调节

ac=[]#为了画图做准备
ls=[]#同上

for epoch in range(Epoch):
    for step, (b_x, b_y) in enumerate(train_loader):
    
        output = cnn(b_x)[0]
        loss = loss_fun(output, b_y)
        optimizer.zero_grad()
        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)
            
            ac.append(accuracy)
            ls.append(loss.data.numpy())

测试

test_output , _= cnn(test_x[:10])
pred_y = torch.max(test_output, 1)[1].data.numpy()
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')
print('准确率',sum(pred_y==test_y[:10].numpy()/10)*100,'%')

画图

ac.append(accuracy)
ls.append(loss.data.numpy())
plt.subplot(121)#把图像分为1行2列，121代表1行2列中的第一个
plt.plot(ls, 'r-',label='loss' )

plt.subplot(122)#把图像分为1行2列，121代表1行2列中的第二个


plt.plot(ac, 'g--',label='ac' )


plt.show()