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

#定义类
#我们通过子类化 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()
print(model)

输出：

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)
  )
)

为了使用这个模型，我们给它传递输入数据。
将数据传递给模型并调用后返回一个 2 维tensor(第0维对应一组 10 个代表每种类型的原始预测值，第1维对应该类型对应的原始预测值)。我们将它传递给一个 nn.Softmax 模块的实例来来获得预测概率。

X = torch.rand(1, 28, 28)
logits = model(X)
pred_probab = nn.Softmax(dim=1)(logits)
#模型的最后一层返回 logits(介于[负无穷,正无穷]之间的原始值)，然后被传递给 nn.Softmax 模块。这些 logits 值被缩放到 [0,1]，代表模型对每种类型的预测概率， dim 参数代表沿着该维度数值应该加总为 1.dim=1是因为数据经过模型训练输出的结果是一个形状为[1, 10]的二维张量。
#dim (int) - Softmax是被计算的维度(当dim=0时，指的是在维度0上的元素相加等于1)。
y_pred = pred_probab.argmax(1)#返回最大值所在索引
print(f"Predicted class: {y_pred}")

输出：
Predicted class: tensor([7])

nn.Sequential 是一个模块的有序容器。数据会沿着模块定义的顺序流动。可以使用 sequential container(译者注：有序容器,也有的书称之为线性容器)来组成一个快速网络。