PyTorch 官方教程：撸一个神经网络

本文为 PyTorch 官方教程中：如何构建神经网络。基于 PyTorch 专门构建神经网络的子模块 torch.nn 构建一个简单的神经网络。

完整教程运行 codelab

torch.nn 文档

神经网络由对数据执行操作的层/模块组成。torch.nn 提供了构建神经网络所需的所有模块。

PyTorch 中的每个模块都是 nn.module 的子类。
在下面的部分中，我们将构建一个神经网络来进行10种类别的分类。

建立神经网络

神经网络由对数据执行操作的层/模块组成。torch.nn 提供了构建神经网络所需的所有模块。PyTorch 中的每个模块都是 nn.module 的子类。
在下面的部分中，我们将构建一个神经网络来进行10种类别的分类。

import os
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

加载训练设备

我们希望能够在硬件加速器，比如 GPU 上训练我们的模型。可以通过 torch.cuda 来检测 GPU 是否可用。

device = 'cuda' if torch.cuda.is_available() else 'cpu' #检测gpu是否可用，不可用使用cpu
print('Using {} device'.format(device)) #输出使用设备类型

定义类

我们通过 nn.Module 来定义神经网络，并在__init__ 中初始化神经网络。每个 nn.Module 子类在 forward 方法中实现对输入数据的操作。

class NeuralNetwork(nn.Module):
    def __init__(self): #定义网络结构
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10),
            nn.ReLU()
        )

    def forward(self, x): #前向传播
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

在使用模型前需要先实例化模型，并将其移动到 GPU 上

model = NeuralNetwork().to(device) #实例化模型
print(model)

为了在模型的输入和输出之间创建复杂的非线性映射，需要使用非线性的激活函数。

它们在线性变换后引入非线性，帮助神经网络学习各种各样的复杂映射。在这个模型中，我们在线性层之间使用 nn.ReLU，也可以使用其他激活函数来引入非线性。

X = torch.rand(1, 28, 28, device=device)  #生成（1，28，28）的数据
logits = model(X) #向模型输入数据
pred_probab = nn.Softmax(dim=1)(logits) #调用softmax 将预测值映射为（0，1）间的概率
y_pred = pred_probab.argmax(1) #最大概率对应分类
print(f"Predicted class: {y_pred}")

神经网络各层说明

接下来，我们分解网络来具体讲述每一层的功能。

为了说明这一点，我们将取小批量的3个尺寸为28x28的图像样本输入网络

input_image = torch.rand(3,28,28) #生成（3，28，28）的数据
print(input_image.size())

nn.Flatten 层

Flatten 层用来把多维的输入一维化，常用在从卷积层到全连接层的过渡。

nn.Flatten 层，可以将每个 28x28 图像转换 784 ($28\times 28=784$)个像素值的连续数组（批量维度保持为3）。

flatten = nn.Flatten() 
flat_image = flatten(input_image) #（3，28，28）转换为（3，784）
print(flat_image.size())

nn.Linear 层

nn.Linear 层，即线性层，是一个使用权重和偏差对输入数据作线性变换的模块。

layer1 = nn.Linear(in_features=28*28, out_features=20) #输入（3，28*28） 输出（3，20）
hidden1 = layer1(flat_image)
print(hidden1.size())

nn.ReLU 层

为了在模型的输入和输出之间创建复杂的非线性映射，需要使用非线性的激活函数。它们在线性变换后引入非线性，帮助神经网络学习各种各样的复杂映射。

在这个模型中，我们在线性层之间使用 nn.ReLU，也可以使用其他激活函数来引入非线性。

print(f"Before ReLU: {hidden1}\n\n")
hidden1 = nn.ReLU()(hidden1)
print(f"After ReLU: {hidden1}")

nn.Sequential 层

神经网络的最后一个线性层返回 logits，即值域区间在$[-\infty ，\infty ]$中的原始值。这些值传递给nn.Softmax模块后，logit被缩放为$[0，1]$区间中，表示模型对每个类的预测概率。

dim参数表示每一维度进行运算的位置，运算结果相加为1。

softmax = nn.Softmax(dim=1)
pred_probab = softmax(logits)

输出模型结构

神经网络中的许多层都是参数化的，即具有相关联的权重和偏差，这些参数在训练中被迭代优化。

子类 nn.Module 自动跟踪模型对象内部定义的所有字段，并使用模型的 parameters() 或 named_parameters() 方法访问所有参数。

我们可以通过模型迭代每个参数，并输出其尺寸和值。

print("Model structure: ", model, "\n\n")

for name, param in model.named_parameters():
    print(f"Layer: {name} | Size: {param.size()} | Values : {param[:2]} \n")

最终输出结果可访问完整教程