Pytorch实现卷积神经网络的分类

本文用于学习Pytorch实现卷积神经网络的分类！！！

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision #数据库
import torch.nn.functional as F

torch.manual_seed(1)

epoches = 1 #训练整批数据的次数
batch_size = 500 #每次训练的样本数
lr = 0.001

#1.获得训练数据集
train_data = torchvision.datasets.MNIST(
    root='./mnist', #下载路径
    train=True, #下载训练数据
    transform=torchvision.transforms.ToTensor(), #将数据转化为tensor类型
    download=False #是否下载MNIST数据集
)
#获得测试数据集
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False
)

#将dataset放入DataLoader中  (batch, channel, 28, 28)
loader = Data.DataLoader(
    dataset=train_data,
    batch_size=batch_size,#设置batch size
    shuffle=True#打乱数据
)

#利用前20个样本和标签测试，并归一化
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[:20]/255
test_y = test_data.test_labels[:20]

#2.前向传播过程
class CNN(nn.Module):
    def __init__(self):
        super(CNN,self).__init__()
        self.conv1=nn.Sequential(#(batch, 1, 28 , 28)
            nn.Conv2d(
                in_channels=1,#输入图片深度
                out_channels=16,#滤波器数量
                kernel_size=5,#滤波器大小
                stride=1,#步长
                padding=2 #保持输出图片大小不变 padding=(kernel_size-stride)/2
            ),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        )#(batch, 16, 14, 14)
        self.conv2=nn.Sequential(#(batch, 16, 14, 14)
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )#(batch, 32, 7, 7)
        self.out = nn.Linear(32*7*7, 10)#全连接层，入口参数是卷积层和输出层的节点数

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)#将(batch, 32, 7, 7)展平为(batch, 32*7*7)
        out = self.out(x)
        return out

# 训练模型的同时保存网络模型参数
def save():
    #3.利用自定义的前向传播过程设计网络
    cnn = CNN()

    #4.设置优化算法、学习率
    optimizer = torch.optim.Adam( cnn.parameters(), lr=lr)

    #5.设置损失函数
    loss_func = torch.nn.CrossEntropyLoss()

    #6.迭代训练
    for epoch in range(epoches):
        for step, (batch_x, batch_y) in enumerate(loader):
            out = cnn(batch_x)#输入训练集，获得当前迭代输出值
            loss = loss_func(out, batch_y)#获得当前迭代的损失

            optimizer.zero_grad()#清除上次迭代的更新梯度
            loss.backward()#反向传播
            optimizer.step()#更新权重

            #打印训练过程中的测试集准确率
            if step%500==0:
                test_out = cnn(test_x)#输入测试集
                #获得当前softmax层最大概率对应的索引值
                pred = torch.max(test_out, 1)[1]
                #将二维压缩为一维
                pred_y = pred.data.numpy().squeeze()
                label_y = test_y.data.numpy()
                accuracy = sum(pred_y == label_y)/test_y.size()
                print("第 %d 个epoch，第 %d 次迭代，准确率为 %.2f"%(epoch+1, step/100+1, accuracy))

    #7.保存模型结构和参数
    torch.save(cnn, 'cnn.pkl')

if __name__ == '__main__':
    save()

参考资料：

1. 莫烦《Pytorch教程》