LeNet训练套路(pytorch官方demo)

目录

1 模型定义

2 训练过程

3 测试过程


1 模型定义

导入包:

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

定义一个继承于nn.Module的类,实现两个方法,在初始化方法中搭建网络层结构,在forward函数中定义正向传播的过程:

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 = 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 = F.relu(self.fc1(x))      # output(120)
        x = F.relu(self.fc2(x))      # output(84)
        x = self.fc3(x)              # output(10)
        return x

定义卷积层的函数Conv2d的函数定义,参数如下:

def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: _size_2_t,
        stride: _size_2_t = 1,
        padding: _size_2_t = 0,
        bias: bool = True,
    ):

in_channels:输入特征矩阵的深度;

out_channels:输出特征矩阵深度,也即使用卷积核的个数;

kernel_size:卷积核大小;

stride:步长,默认为1;

padding:在四周进行补0处理,默认为0;

bias:偏置,默认为True,即使用偏置。

定义池化层的函数MaxPool2d的函数定义,参数如下:

def __init__(
        self, 
        kernel_size: _size_any_t, 
        stride: Optional[_size_any_t] = None,
        padding: _size_any_t = 0
    ) -> None:

kernel_size:池化核大小;

stride:步长,若不指定会采用和池化核大小一样的步长。


输入为3×32×32,经过conv1后根据公式:

N=(W-F+2P)/S+1

计算出输出的矩阵是28×28的,由于使用了16个卷积核,因此输出为16×28×28;

再经过pool1,池化层不改变深度,利用2×2的池化核得到的输出为16×14×14;

经过conv2,此时输入的深度是16,采用32个卷积核,大小为5×5,同理根据公式计算,得到的输出为32×10×10;

再经过pool2,利用2×2的池化核得到的输出为32×5×5;

经过fc1,全连接层的输入为一维向量,要将输入展平,节点设置根据网络结构设为120;

经过fc2,输入设置为120,节点设置根据网络结构设为84;

经过fc3,输入设置为84,节点设置根据输出的类别,设为10。

网络结构定义好之后进行前向传播,经过卷积层后要通过一个relu激活函数;在进入全连接层之前要将特征矩阵展平,其中view函数的第一个参数设为1,表示动态调整这个维度上的元素个数,以保证元素的总数不变;经过全连接层需要通过relu激活函数,但在最后一层按理说要接一个softmax层,将输出转化为一个概率分布,但在训练网络过程中计算卷积交叉熵时,已经在内部实现了一个更加高效的softmax方法,所以不需要添加softmax层了。


2 训练过程

采用CIFAR10数据集,一共有10个类别,图像均为3×32×32的,首先从网络上下载数据集:

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

将train设为True会下载数据集中的训练集样本,将download改为True会下载数据集到root所指定的文件夹内,transform是对图像进行预处理的函数:

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

ToTensor函数所做的事情是将输入图像的高度×宽度×深度改为深度×高度×宽度,将每个维度的像素值范围从[0, 255]改为[0, 1];Normalize函数是标准化的过程。

然后将数据分成不同批次:

train_loader = torch.utils.data.DataLoader(train_set, batch_size=36,
                                           shuffle=True, num_workers=0)

batch_size即为每一批次的数据量,shuffle设为True可以将数据打乱。

测试集的导入方式同理,测试集导入之后,将其转化为一个迭代器,然后通过next方法可以获取到一批数据,包括图像以及对应标签:

val_data_iter = iter(val_loader)
val_image, val_label = val_data_iter.next()

然后将模型实例化,定义损失函数以及优化器:

net = LeNet()
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)

其中CrossEntropyLoss函数中包括了softmax函数,因此全连接层最后一层不需要加softmax层;优化器使用Adam优化器,第一个参数表示训练net中的所有可训练的参数(parameters)都进行训练,学习率设为0.001。

然后进入训练过程:

for epoch in range(5):  # 循环5次数据集

    running_loss = 0.0  # 该变量用来累加训练过程中的损失
    for step, data in enumerate(train_loader, start=0):  # start=0代表step从0开始
        # enumerate获取输入;data是[输入,标签]的列表
        inputs, labels = data

        # 将参数梯度归零,否则梯度会一直累加
        optimizer.zero_grad()
        # 前向传播 + 反向传播 + 优化
        # 将输入送到net中进行前向传播得到输出
        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]
                # 找到概率最大的值是哪一个,要在维度1上找最大值,第0维是batch,[1]是只需要知道索引值即可
                predict_y = torch.max(outputs, dim=1)[1]
                # 将预测值与实际值作比较。相同返回1,不同返回0,通过求和函数可知本次预测对了多少个样本
                # 此时是ToTensor类型,通过item获取数值,除以测试样本数目,得到准确率
                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清零
                running_loss = 0.0

训练完成后保存模型:

# 保存模型权重
save_path = './Lenet.pth'
torch.save(net.state_dict(), save_path)

3 测试过程

训练完成后可以自己搜索图片输入进行预测:

# 由于输入图片大小不确定,因此多了一步将图片缩放成32×32的大小
transform = transforms.Compose(
    [transforms.Resize((32, 32)),
     transforms.ToTensor(),
     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()
# 载入权重文件
net.load_state_dict(torch.load('Lenet.pth'))

# 载入图像
im = Image.open('1.jpg')
im = transform(im)  # [C, H, W]
# 由于输入的图像应该是4个维度,因此在前面多加1个维度
im = torch.unsqueeze(im, dim=0)  # [N, C, H, W]

with torch.no_grad():
    outputs = net(im)
    predict = torch.max(outputs, dim=1)[1].numpy()
# 利用最大值的index找到所对应的类
print(classes[int(predict)])

猜你喜欢

转载自blog.csdn.net/AdjsWsgz/article/details/127313701