神经网络包括在数据上执行操作的layers/modules。torch.nn命名空间提供建立你自己的神经网络模块。每个模块继承nn.Module。一个神经网络是一个由其他模块（层）组成的模块。这种嵌套的结构考虑到容易建立和管理复杂的架构。

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

获取用于训练的设备

在硬件加速器（GPU或MPS）上训练模型，检查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")

定义模型类

通过继承nn.Mudule类，定义神经网络模型类。and在__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

创建神经网络实例，并将它移动到device上，并打印其结构。

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

为了使用模型，我们将输入数据传递给它。这样执行模型的forward方法，随之执行背后的操作，并不是直接执行model.forward()。

调用模型，返回2维tensor，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}")

模型的层

创建3张28X28大小的图片，看经过每层神经网络发生了什么。