4. BUILD THE NEURAL NETWORK

4. BUILD THE NEURAL NETWORK

神经网络由对数据执行操作的层/模块组成。torch.nn提供了构建自己的神经网络所需的所有构建块。PyTorch中的每个模块都是nn.module的子类。神经网络是由其他模块（层）组成的模块本身。这种嵌套结构允许轻松地构建和管理复杂的体系结构。

接下来，我们将构建一个神经网络来分类FashionMNIST数据集中的图片。

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

4.1 获取训练设备

检查torch.cuda或torch.backends.mps是否可用，否则使用CPU进行训练。

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

4.2 定义类

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

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__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),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

创建一个NeuralNetwork的实例，将其移动到设备上并打印其结构

model = NeuralNetwork().to(device)
print(model)

为了使用该模型，我们将输入数据传递给它。这将执行模型的forward以及一些后台操作。
对输入调用模型返回二维张量，dim=0对应于每个类的10个原始预测值的每个输出，dim=1对应于每个输出的各个值。通过nn.Softmax模块的实例来获得预测概率。

X = torch.rand(1, 28, 28, device=device)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
y_pred = pred_probab.argmax(1)
print(f"Predicted class: {y_pred}")

4.3 模型层级

接下来我将分解FashionMNIST模型中的层次。为了更好地进行说明，我们将取出一个由3张尺寸为28x28的图像组成的迷你批次样本，看看通过网络时会发生什么。

input_image = torch.rand(3,28,28)
print(f"input_image size:{input_image.size()}")

input_image size:torch.Size([3, 28, 28])

nn.Flatten

初始化nn.Flatten层，以将每个2D 28x28图像转换为784个像素值的连续阵列（保持迷你批量维度（dim=0））。

flatten = nn.Flatten()
flat_image = flatten(input_image)
print(f"flat_image size:{flat_image.size()}")

flat_image size:torch.Size([3, 784])

nn.Linear

线性层是使用其存储的权重和偏差对输入应用线性变换的模块。

layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(f"hidden1 size:{hidden1.size()}")

hidden1 size:torch.Size([3, 20])

nn.ReLU

非线性激活是在模型的输入和输出之间创建复杂映射的原因。它们在线性变换后应用，引入非线性，帮助神经网络学习各种各样的现象。
在这个模型中，我们在线性层之间使用nn.ReLU，但还有其他激活会在模型中引入非线性。

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

Before ReLU: tensor([[-0.2347, -0.1786, -0.3838, 0.2101, -0.0453, 0.2431, 0.6029, -0.5038,
-0.0544, 0.0120, -0.2598, 0.1301, 0.2343, -0.3681, 0.0421, 0.1453,
-0.5293, 0.1022, -0.3209, -0.6937],
[ 0.1206, -0.3236, 0.0152, 0.3843, 0.1214, 0.2459, 0.8221, -0.4659,
0.0657, -0.4249, -0.4321, 0.2330, -0.2569, -0.3014, 0.0575, 0.0606,
-0.3889, 0.1250, -0.1252, -0.5308],
[-0.2399, -0.1596, 0.0495, 0.7113, 0.1518, 0.0298, 0.6635, -0.4252,
-0.0991, -0.3769, -0.6458, 0.0348, -0.0304, -0.3883, 0.0419, 0.2834,
-0.3759, -0.0255, -0.1007, -0.3818]], grad_fn=)

After ReLU: tensor([[0.0000, 0.0000, 0.0000, 0.2101, 0.0000, 0.2431, 0.6029, 0.0000, 0.0000,
0.0120, 0.0000, 0.1301, 0.2343, 0.0000, 0.0421, 0.1453, 0.0000, 0.1022,
0.0000, 0.0000],
[0.1206, 0.0000, 0.0152, 0.3843, 0.1214, 0.2459, 0.8221, 0.0000, 0.0657,
0.0000, 0.0000, 0.2330, 0.0000, 0.0000, 0.0575, 0.0606, 0.0000, 0.1250,
0.0000, 0.0000],
[0.0000, 0.0000, 0.0495, 0.7113, 0.1518, 0.0298, 0.6635, 0.0000, 0.0000,
0.0000, 0.0000, 0.0348, 0.0000, 0.0000, 0.0419, 0.2834, 0.0000, 0.0000,
0.0000, 0.0000]], grad_fn=)

nn.Sequential

nn.Sequential是模块的有序容器。数据以定义的相同顺序通过所有模块。您可以使用顺序容器将类似seq_modules的快速网络组合在一起。

seq_modules = nn.Sequential(
    flatten,
    layer1,
    nn.ReLU(),
    nn.Linear(20, 10)
)
input_image = torch.rand(3,28,28)
logits = seq_modules(input_image)

4.4 模型参数

神经网络内部的许多层都是参数化的，即具有在训练期间优化的相关权重和偏差。子类nn.Module自动跟踪模型对象中定义的所有字段，并使用模型的parameters()或named_parameters()方法访问所有参数。

在下面的代码里，对每个参数进行遍历，并打印。

print(f"Model structure: {model}\n\n")

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

Model structure: NeuralNetwork(
(flatten): Flatten(start_dim=1, end_dim=-1)
(linear_relu_stack): Sequential(
(0): Linear(in_features=784, out_features=512, bias=True)
(1): ReLU()
(2): Linear(in_features=512, out_features=512, bias=True)
(3): ReLU()
(4): Linear(in_features=512, out_features=10, bias=True)
)
)

Layer: linear_relu_stack.0.weight | Size: torch.Size([512, 784]) | Values : tensor([[ 0.0328, 0.0088, -0.0321, …, -0.0181, -0.0218, -0.0153],
[-0.0059, 0.0152, 0.0132, …, 0.0041, -0.0064, -0.0010]],
device=‘cuda:0’, grad_fn=)

Layer: linear_relu_stack.0.bias | Size: torch.Size([512]) | Values : tensor([-0.0144, -0.0142], device=‘cuda:0’, grad_fn=)

Layer: linear_relu_stack.2.weight | Size: torch.Size([512, 512]) | Values : tensor([[-0.0281, 0.0108, 0.0422, …, 0.0082, -0.0360, 0.0323],
[ 0.0253, -0.0117, -0.0169, …, 0.0260, -0.0238, 0.0284]],
device=‘cuda:0’, grad_fn=)

Layer: linear_relu_stack.2.bias | Size: torch.Size([512]) | Values : tensor([-0.0028, 0.0026], device=‘cuda:0’, grad_fn=)

Layer: linear_relu_stack.4.weight | Size: torch.Size([10, 512]) | Values : tensor([[ 0.0380, 0.0085, 0.0147, …, -0.0246, 0.0411, 0.0124],
[-0.0414, -0.0192, 0.0350, …, -0.0145, -0.0048, 0.0338]],
device=‘cuda:0’, grad_fn=)

Layer: linear_relu_stack.4.bias | Size: torch.Size([10]) | Values : tensor([-0.0010, -0.0017], device=‘cuda:0’, grad_fn=)

参考资料

1. https://pytorch.org/tutorials/beginner/basics/buildmodel_tutorial.html
2. https://blog.csdn.net/weixin_43135178/article/details/114602960
3. https://www.bilibili.com/video/BV1MT411p7CZ