基于PyTorch学AI——构建网络模型

有了前面的铺垫，终于可以利用PyTorch构建一个神经网络了。
在这里插入图片描述

神经网络由很多节点组成，每个节点有输入输出，所有节点组成一个网络结构，整体网络也有输入和输出，通过输入训练数据，输出标签，完成一个特定任务。
PyTorch中利用nn.Module类实现节点功能。

构建神经网络

先说说基本逻辑，神经网络由层和模块组成，可以对数据执行某种操作。
torch.nn命名空间提供了构建一个神经网络所需的各种模块，每个模块都继承自nn.Module，一个模块可以包含多个其它模块，利用这种嵌套结构可以轻松构建出复杂的神经网络模型。

下面就开始构建FashionMNIST数据集的图片分类模型。
先导入相关库，最重要的是nn库。

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

指定训练设备

如果有GPU，尽量使用GPU训练模型。
利用torch.cuda.is_available()方法检测cuda是否可用，如果不可用，则使用CPU。

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

定义模型网络

自定义的神经网络模型都继承自nn.Module，需要实现init方法和forward方法。
__init__方法中定义其子模块，构成整个网络。
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模块，有多个子模块组成。

• flatten 子模块
  用于把一个tensor”压扁“。看例子秒懂：

>>> input = torch.randn(32, 1, 5, 5)
>>> # With default parameters
&