使用pytorch搭建LeNet网络 学习笔记

视频教程来自B站up@霹雳吧啦Wz，链接:2.1 pytorch官方demo(Lenet)

 

model

 

model.py

import torch.nn as nn
import torch.nn.functional as F


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, 5)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(16, 32, 5)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(32*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):            #x为输入的数据
        x = F.relu(self.conv1(x))    # input(3, 32, 32) output(16, 28, 28)
        x = self.pool1(x)            # output(16, 14, 14)
        x = F.relu(self.conv2(x))    # output(32, 10, 10)
        x = self.pool2(x)            # output(32, 5, 5)
        x = x.view(-1, 32*5*5)       # output(32*5*5)  x.view()函数将x展成一维向量
        x = F.relu(self.fc1(x))      # output(120)
        x = F.relu(self.fc2(x))      # output(84)
        x = self.fc3(x)              # output(10)  
        return x
        

 
首先定义一个类LeNet，继承于父类nn.module。LeNet中定义两个方法：初始化函数（搭建网络层结构）和定义正向传播过程。
 

nn.Conv2d函数

self.conv1 = nn.Conv2d(3, 16, 5)  #定义第一个卷积层

# nn.Conv2d 为定义卷积层函数,参数如下：

def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1,
                 bias=True, padding_mode='zeros'):
# in_channels（int）–输入图像中的通道数
# out_channels（int）–卷积产生的通道数（卷积核个数）
# kernel_size（int或tuple）–卷积核的大小
# stride（int或tuple，可选）–卷积的步幅。默认值：1
# padding（int或tuple，可选）–将零填充添加到输入的两侧。默认值：0
# 				  若为int型，比如1，则在上下左右都补一行(列)零
# 			      若为tuple:(1,2) 1代表上下方各补一行零，2代表左右两侧各补两列零
# 其他参数用的较少，参考https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html?highlight=conv2d#torch.nn.Conv2d

 

nn.MaxPool2d函数

self.pool1 = nn.MaxPool2d(2, 2)   #定义一个下采样层 

def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
                 return_indices=False, ceil_mode=False):

# kernel_size –取最大值的窗口的大小
# stride –窗口的步幅,默认值为kernel_size
# padding –在两边都添加隐式零填充

 
根据下图网络结构和公式，计算出每层的参数

 

train

 

train.py

import torch
import torchvision
import torch.nn as nn
from model import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np


def main():
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # 50000张训练图片
    # 第一次使用时要将download设置为True才会自动去下载数据集
    train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                             download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
                                               shuffle=True, num_workers=0)

    # 10000张验证图片
    # 第一次使用时要将download设置为True才会自动去下载数据集
    val_set = torchvision.datasets.CIFAR10(root='./data', train=False,
                                           download=False, transform=transform) 
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=4,
                                             shuffle=False, num_workers=0) 
                                           
    # 获得图像和标签值
    val_data_iter = iter(val_loader)
    val_image, val_label = val_data_iter.next()

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    # 交叉熵损失函数，该损失函数结合了nn.LogSoftmax()和nn.NLLLoss()两个函数，分类问题经常使用
    # Adam优化器，第一个参数为所需要训练的参数
    net = LeNet()
    loss_function = nn.CrossEntropyLoss()                 #定义损失函数
    optimizer = optim.Adam(net.parameters(), lr=0.001)    #定义优化器，parameters()所有可训练参数都进行训练
                                                          #lr 学习率

    for epoch in range(5):  # 训练集迭代五次

        running_loss = 0.0  #累加损失
        #遍历训练集样本，从零开始，
        for step, data in enumerate(train_loader, start=0):  #遍历训练集样本，从零开始
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients(损失梯度清零)，不清零能变相增加batch_size
            optimizer.zero_grad()
            # forward + backward + optimize
            outputs = net(inputs)                   
            loss = loss_function(outputs, labels)   #计算损失(网络预测值，真实标签)
            loss.backward()                         #反向传播
            optimizer.step()                        #参数更新

            # print statistics
            running_loss += loss.item()
            if step % 500 == 499:        # 每训练500步打印：
                with torch.no_grad():    # 接下来计算过程中不要计算每个节点的误差损失梯度(节省内存)
                    outputs = net(val_image)  # [batch, 10]
                    predict_y = torch.max(outputs, dim=1)[1]   # 在输出第1个维度上寻找最大值，第0个维度为batch
                    accuracy = torch.eq(predict_y, val_label).sum().item() / val_label.size(0)

                    print('[%d, %5d] train_loss: %.3f  test_accuracy: %.3f' %
                          (epoch + 1, step + 1, running_loss / 500, accuracy))
                    running_loss = 0.0

    print('Finished Training')

    save_path = './Lenet.pth'
    torch.save(net.state_dict(), save_path)


if __name__ == '__main__':
    main()

 

下载数据集

train_set = torchvision.datasets.CIFAR10(root='./data', train=True,
                                             download=True, transform=transform)

# root='./data'        将数据集下载到当前目录data文件夹下
# train=True           为Treu时将导入训练集样本
# download=True        为Treu时将下载数据集，下载完改为False
# transform=transform  预处理函数

 

导入数据集

train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
                                               shuffle=True, num_workers=0)
# rain_set        下载好的数据集
# batch_size=36   每一批随机拿出36张图片进行训练
# shuffle=True    是否打乱

 

predict

下载一张飞机图片，放入当前目录下，保存为1.png

 

predict.py

import torch
import torchvision.transforms as transforms
from PIL import Image

from model import LeNet


def main():
    transform = transforms.Compose(
        [transforms.Resize((32, 32)),		#将图片缩放为32*32大小
         transforms.ToTensor(),				#转化为Tensor
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])	#标准化处理

    classes = ('plane', 'car', 'bird', 'cat',
               'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    net = LeNet()   								#实例化LeNet
    net.load_state_dict(torch.load('Lenet.pth'))	#载入权重文件

    im = Image.open('1.png')  						# 导入图片
    im = transform(im)  # [C, H, W]					# 预处理为标准图片大小(32*32)
    im = torch.unsqueeze(im, dim=0)  # [N, C, H, W] # 加一个新的维度

    with torch.no_grad():
        outputs = net(im)
        predict = torch.max(outputs, dim=1)[1].data.numpy()
    print(classes[int(predict)])


if __name__ == '__main__':
    main()

 
运行后，输出plane