Pytorch搭建自定义神经网络

简介

在日常的实验中，为了改变网络模型，我们一般会自己搭建网络，在基础上进行修改创新，本文主要介绍如何利用Pytorch搭建自定义的神经网络

nn.Module类

在Pytorch中搭建自己的网络模型，首先要继承nn.Module，通过继承的方式定义自己的模型。

init函数

在Python中__init__函数就是构造函数，在__init__函数中我们一般会定义模型中的一些结构，比如卷积层，池化层，全连接层等。这里会用到一个函数nn.Sequential，这个函数的目的是为了拼接各种卷积，全连接等（当然你也可以把里面的卷积等结构拆开写）。

forward函数

forward函数是重写了nn.Module中的forward函数，目的在于训练/预测数据，forward函数中有个参数x表示输入的数据，如果__init__函数中定义的参数没有问题，那么在forward中直接调用就可以，最后返回输出的结果即可。
上面的文字理解或许不是很好理解，下面我们看一下利用Pytorch搭建的AlextNet网络模型。

import torch.nn as nn


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        # 卷积层
        self.layers = nn.Sequential(
            # 第一层
            nn.Conv2d(3, 96, kernel_size=(3, 3)),
            nn.BatchNorm2d(96),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=3, stride=2),

            # 第二层
            nn.Conv2d(96, 256, kernel_size=(3, 3)),
            nn.BatchNorm2d(256),
            nn.ReLU(True),
            nn.MaxPool2d(kernel_size=3, stride=2),

            # 第三层
            nn.Conv2d(256, 384, kernel_size=(3, 3), padding=1),
            nn.ReLU(True),

            # 第四层
            nn.Conv2d(384, 384, kernel_size=(3, 3), padding=1),
            nn.ReLU(True),

            # 第五层
            nn.Conv2d(384, 256, kernel_size=(3, 3), padding=1),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )

        # 全连接层
        self.fc = nn.Sequential(
            nn.Linear(1024, 2048),
            nn.Dropout(0.5),
            nn.Linear(2048, 2048),
            nn.Dropout(0.5),
            nn.Linear(2048, 10)
        )

    def forward(self, x):
        x = self.layers(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

分析

根据上面的介绍，首先继承nn.Module，然后编写__init__函数和forward函数。这里我详细的分析一下每行代码的意义。

• nn.Conv2d：表示一个2D卷积层，nn.Conv2d(3, 96, kernel_size=(3, 3))表示这是一个2D卷积层，其中输入的是3通道，输出的是96通道，卷积核大小是(3,3)
• nn.MaxPool2d：表示一个2D最大池化层，nn.MaxPool2d(kernel_size=3, stride=2)表示池化层大小为(3,3)，步长为2（注意：池化层不改变输入的维度）
• nn.ReLU：表示使用relu作为激活函数
• nn.BatchNorm2d：表示批归一化处理
• nn.Linear()：全连接层，nn.Linear(2048, 10)表示输入为2048输出为10
• nn.Dropout()：nn.Dropout(0.5)表示舍弃50%的参数

利用nn.Sequential把这些结构进行整合，最后在forward函数中进行调用，返回处理后的值即可。因此利用Pytorch搭建网络可以分为以下几步：

1. 定义网络名称继承nn.Module
2. 在构造方法中定义需要的结构
3. 在forward函数中调用，并返回处理后的值

预测

利用刚才建立的网络模型在cifar10数据集上进行测试，加载数据集的方式不再赘述，下面主要介绍一下如何训练+预测。

1. 生成已经建立的网络模型并利用GPU加速：alexNet = AlexNet().to(device)
2. 定义优化器optimize = torch.optim.Adam(alexNet.parameters(), lr=0.01)
3. 定义损失函数loss_function = nn.CrossEntropyLoss()
4. 定义迭代次数开始预测
5. 反向传播+优化器优化

预测代码如下：

import os

import torch
import torch.nn as nn
from torchvision import transforms
from torchvision import datasets
import torch.utils.data
from AlexNet import AlexNet

device = torch.device("cuda")

if __name__ == '__main__':
    download = False

    if 'res' not in os.listdir():
        download = True

    train_set = datasets.CIFAR10(root='./res/data', transform=transforms.ToTensor(), train=True, download=download)
    test_set = datasets.CIFAR10(root='./res/data', transform=transforms.ToTensor(), train=False, download=download)
    train_set = torch.utils.data.DataLoader(dataset=train_set, batch_size=64, shuffle=True, num_workers=4)

    test_set = torch.utils.data.DataLoader(dataset=test_set, batch_size=64, shuffle=True, num_workers=4)
    alexNet = AlexNet().to(device)

    optimize = torch.optim.Adam(alexNet.parameters(), lr=0.01)
    loss_function = nn.CrossEntropyLoss()
    epochs = 5
    for epoch in range(epochs):
        loss_sum = 0.0
        for i, data in enumerate(train_set):
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            optimize.zero_grad()
            outputs = alexNet(inputs)
            loss_per = loss_function(outputs, labels)
            loss_per.backward()
            optimize.step()
            loss_sum += loss_per.item()

            if i % 100 == 99:
                print('[Epoch:%d, batch:%d] train loss: %.03f' % (epoch + 1, i + 1, loss_sum / 100))
                loss_sum = 0.0

        total, right = 0, 0
        for i, data in enumerate(test_set):
            test_inputs, test_labels = data
            test_inputs, test_labels = test_inputs.to(device), test_labels.to(device)
            test_outputs = alexNet(test_inputs)
            test_outputs = torch.max(test_outputs.data, 1)[1]
            total += test_outputs.size(0)
            right += (test_labels == test_outputs).sum()
        print("第{}轮的准确率为:{:.2f}%".format(epoch + 1, 100.0 * right.item() / total))