PyTorch 教程-快速上手指南

文章目录

PyTorch Quickstart
- 1.处理数据
- 2.创建模型
- 3.优化模型参数
- 4.保存模型
- 5.加载模型
PyTorch 基础入门
- 1.Tensors
- - 1.1初始化张量
  - 1.2张量的属性
  - 1.3张量运算
  - - 1.3.1张量的索引和切片
    - 1.3.2张量的连接
    - 1.3.3算术运算
    - 1.3.4单元素张量转变为Python数值
  - 1.4Tensor与NumPy的桥接
  - - 1.4.1Tensor to NumPy array
    - 1.4.2NumPy array to Tensor
- 2.在PyTorch中加载数据集
- - 2.1装载数据集
  - 2.2迭代和可视化数据集
  - 2.3为文件创建自定义数据集
  - 2.4使用DataLoader为训练准备数据
  - 2.5遍历数据加载器
- 3.Transforms
- 4.构建一个神经网络
- - 4.1导入包
  - 4.2检查GPU是否可用
  - 4.3定义类
  - 4.4模型层
  - - 4.4.1nn.Flatten
    - 4.4.2nn.Linear
    - 4.4.3nn.ReLU
    - 4.4.4nn.Sequential
    - 4.4.5nn.Softmax
  - 4.5模型参数
- 5.使用 torch.autograd 进行自动微分
- - 5.1张量，函数与计算图
  - 5.2计算梯度
  - 5.3禁用梯度跟踪
  - 5.4计算图更多信息
  - 5.5张量梯度和雅可比积（可选）
- 6.优化模型参数
- - 6.1前提代码
  - 6.2超参数
  - 6.3优化循环
  - 6.4损失函数
  - 6.5优化器
  - 6.6完整实现
- 7.模型保存和加载
- - 7.1模型权重的保存和加载
  - 7.2保存和加载模型结构
参考

说明：本教程翻译自 Pytorch 官方教程： Introduction to PyTorch ，适合对 Python 和深度学习具有一定基础的同学学习，是 Pytorch 的入门教程。 😃

PyTorch Quickstart

1.处理数据

PyTorch有两个处理数据的基本操作：torch.utils.data.DataLoader和torch.utils.data.Dataset。Dataset用于存储样本及其对应的标签，而DataLoader则围绕Dataset包装了一个可迭代的数据加载器。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

PyTorch 提供特定于领域的库，如 TorchText, TorchVision 和 TorchAudio，所有这些库都包含数据集。对于本教程，将使用 TorchVision 数据集。

torchvision.datasets模块包含了许多真实世界视觉数据的 Dataset对象，比如 CIFAR、 COCO (完整列表在这里)。在本教程中，我们使用 FashionMNIST 数据集。每个 TorchVision Dataset都包含两个参数: transform 和 target_transform，分别用于转换样本和标签。

# 从开源数据集下载训练数据。
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

# 从开源数据集下载测试数据。
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)

将Dataset作为参数传递给DataLoader。这将在数据集上包装一个迭代器，并支持自动批处理、采样、随机打乱和多进程数据加载。在这里，定义了一个大小为64的批处理**，即 DataLoader 迭代器中的每个元素都会返回一个由64个特征和标签组成的批次数据**。

batch_size = 64

# 创建数据加载器
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

for X, y in test_dataloader:
    # N 表示批量大小（batch size），C 表示通道数（channels），H 表示图像高度（height），W 表示图像宽度（width）。
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

更详细的内容请查看 loading data in PyTorch 。

2.创建模型

为了在 PyTorch 中定义一个神经网络，需要创建一个继承自 nn.Module 的自定义类。在 __init__ 方法中定义网络的层次结构，并在 forward 方法中指定数据将如何通过网络的各个层。为了加速神经网络中的操作，我们将其移动到 GPU 或 MPS (如果有的话)。

# 获取 cpu, gpu 或 mps 设备用于加速训练.
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

# 定义神经网络
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

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

更详细的内容请查看 building neural networks in PyTorch 。

3.优化模型参数

为了训练一个模型，我们需要一个 loss function 和一个optimizer 。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

在单个训练循环中，模型对训练数据集（以批次 batch 输入）进行预测，并反向传播预测误差以调整模型的参数。

def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # 将数据移动到 GPU 上
        X, y = X.to(device), y.to(device)

        # 计算预测值与损失
        pred = model(X)
        loss = loss_fn(pred, y)

        # 反向传播
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        if batch % 100 == 0:
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")

我们还针对测试数据集检查模型的性能，以确保它正在学习。

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num_batches
    correct /= size
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

训练过程在多个迭代（epoch）中进行。在每个 epoch 中，模型学习参数以做出更好的预测。我们在每个 epoch 打印模型的准确度和损失；我们希望看到准确度随着每个 epoch 的增加而增加，损失随着每个 epoch 的增加而减少。

epochs = 5
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train(train_dataloader, model, loss_fn, optimizer)
    test(test_dataloader, model, loss_fn)
print("Done!")

更详细的内容请查看 Training your model 。

4.保存模型

保存模型的常见方法是序列化内部状态字典（包含模型参数）。

torch.save(model.state_dict(), "model.pth")
print("Saved PyTorch Model State to model.pth")

5.加载模型

加载模型的过程包括重新创建模型结构并将状态字典加载到其中。

model = NeuralNetwork().to(device)
model.load_state_dict(torch.load("model.pth"))

这个模型现在可以用来做预测。

classes = [
    "T-shirt/top",
    "Trouser",
    "Pullover",
    "Dress",
    "Coat",
    "Sandal",
    "Shirt",
    "Sneaker",
    "Bag",
    "Ankle boot",
]

model.eval()
x, y = test_data[0][0], test_data[0][1]
with torch.no_grad():
    x = x.to(device)
    pred = model(x)
    predicted, actual = classes[pred[0].argmax(0)], classes[y]
    print(f'Predicted: "{predicted}", Actual: "{actual}"')

更详细的内容请查看 Saving & Loading your model 。

PyTorch 基础入门

1.Tensors

张量是一种专门的数据结构，非常类似于数组和矩阵。在PyTorch中，我们使用张量来编码模型的输入和输出，以及模型的参数。张量类似于NumPy的ndarrays，唯一的区别在于张量可以在GPU或其他硬件加速器上运行。实际上，张量和NumPy数组通常可以共享相同的底层内存，消除了复制数据的需要。张量还针对自动微分进行了优化。

import torch
import numpy as np

1.1初始化张量

张量可以用各种方式初始化。请看下面的例子：

Directly from data

张量可以直接从数据中创建。数据类型是自动推断的。

data = [[1, 2],[3, 4]]
x_data = torch.tensor(data)

From a NumPy array

张量可以从NumPy数组中创建（反之亦然——参见NumPy桥接）。

np_array = np.array(data)
x_np = torch.from_numpy(np_array)

From another tensor

新张量保留参数张量的属性（形状、数据类型），除非被显式覆盖。

x_ones = torch.ones_like(x_data) # retains the properties of x_data
print(f"Ones Tensor: \n {x_ones} \n")

x_rand = torch.rand_like(x_data, dtype=torch.float) # overrides the datatype of x_data
print(f"Random Tensor: \n {x_rand} \n")

Out:

With random or constant values

shape是张量维度的元组。在下面的函数中，它决定了输出张量的维数。

shape = (2,3,)
rand_tensor = torch.rand(shape)
ones_tensor = torch.ones(shape)
zeros_tensor = torch.zeros(shape)

print(f"Random Tensor: \n {rand_tensor} \n")
print(f"Ones Tensor: \n {ones_tensor} \n")
print(f"Zeros Tensor: \n {zeros_tensor}")

Out:

1.2张量的属性

张量属性描述了它们的形状（shape）、数据类型（datatype）和存储它们的设备。

tensor = torch.rand(3,4)

print(f"Shape of tensor: {tensor.shape}")
print(f"Datatype of tensor: {tensor.dtype}")
print(f"Device tensor is stored on: {tensor.device}")

Out:

1.3张量运算

在这里，详细描述了超过100种张量操作，包括算术、线性代数、矩阵操作（转置、索引、切片）、采样等。每个操作都可以在GPU上运行（通常比在CPU上运行速度更快）。

默认情况下，张量是在CPU上创建的。需要使用.to方法将张量明确地移动到GPU上（在检查GPU是否可用后）。请注意，跨设备复制大型张量可能会在时间和内存方面产生昂贵的开销！

# We move our tensor to the GPU if available
if torch.cuda.is_available():
  tensor = tensor.to('cuda')

尝试列表中的一些操作。如果你熟悉NumPy API，你会发现张量API使用起来轻而易举。

1.3.1张量的索引和切片

tensor = torch.ones(4, 4)
print('First row: ',tensor[0])
print('First column: ', tensor[:, 0])
print('Last column:', tensor[..., -1])
tensor[:,1] = 0
print(tensor)

Out:

1.3.2张量的连接

你可以使用torch.cat沿着给定的维度连接一系列张量。另外，你还可以了解torch.stack，它是另一种张量拼接操作，与torch.cat有一些微妙的区别。

t1 = torch.cat([tensor, tensor, tensor], dim=1)
print(t1)

Out:

1.3.3算术运算

# 这计算两个张量之间的矩阵乘法。y1，y2，y3将具有相同的值
y1 = tensor @ tensor.T
y2 = tensor.matmul(tensor.T)

y3 = torch.rand_like(tensor)
torch.matmul(tensor, tensor.T, out=y3)


# 这将计算元素乘积。z1，z2，z3将具有相同的值
z1 = tensor * tensor
z2 = tensor.mul(tensor)

z3 = torch.rand_like(tensor)
torch.mul(tensor, tensor, out=z3)

1.3.4单元素张量转变为Python数值

如果你有一个只包含一个元素的张量，例如将张量中的所有值聚合成一个值，你可以使用 item() 方法将其转换为 Python 数值。

agg = tensor.sum()
agg_item = agg.item()
print(agg_item, type(agg_item))

Out:

In-place操作

In-place操作是将结果存储到操作数中的操作。它们用 _ 后缀表示。例如：x.copy_(y)，x.t_()，将改变 x。

print(tensor, "\n")
tensor.add_(5)
print(tensor)

Out:

【注：In-place操作节省了一些内存，但在计算导数时可能会有问题，因为会立即丢失历史记录。因此，不鼓励使用它们。】

1.4Tensor与NumPy的桥接

Tensor和NumPy之间进行数据交换的机制：在 CPU 上的张量和 NumPy 数组可以共享它们的底层内存位置，改变其中一个将会改变另一个。

1.4.1Tensor to NumPy array

t = torch.ones(5)
print(f"t: {t}")
n = t.numpy()
print(f"n: {n}")

Out:

张量的变化反映在NumPy数组中。

t.add_(1)
print(f"t: {t}")
print(f"n: {n}")

Out:

1.4.2NumPy array to Tensor

n = np.ones(5)
t = torch.from_numpy(n)

NumPy数组的变化反映在张量中。

np.add(n, 1, out=n)
print(f"t: {t}")
print(f"n: {n}")

Out:

2.在PyTorch中加载数据集

处理数据样本的代码可能会变得混乱且难以维护；理想情况下，我们希望数据集代码与模型训练代码解耦，以提高可读性和模块化性。PyTorch 提供了两个数据原语：torch.utils.data.DataLoader 和 torch.utils.data.Dataset，它们允许你使用预加载的数据集以及你自己的数据。Dataset 存储样本及其对应的标签，而 DataLoader 则在 Dataset 周围包装了一个可迭代对象，以便轻松访问样本。

PyTorch 领域库提供了许多预加载的数据集（例如 FashionMNIST），它们是 torch.utils.data.Dataset 的子类，并实现了特定于特定数据的功能。你可以用它们来原型设计和评估模型性能。您可以在这里找到这些数据集：图像数据集、文本数据集和音频数据集。

2.1装载数据集

这是一个如何从TorchVision加载Fashion-MNIST数据集的示例。Fashion-MNIST是由Zalando公司提供的包含6万个训练样本和1万个测试样本的服装图像数据集。每个样本包含一个28x28的灰度图像和对应的10个类别之一的标签。

我们使用以下参数加载FashionMNIST Dataset:

root是存储训练/测试数据的路径,
train指定是训练数据集还是测试数据集,
download=True如果在root路径中不存在数据,则从网络下载数据。
transform和target_transform分别指定对特征和标签的转换。

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt


training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

Out:

2.2迭代和可视化数据集

我们可以像处理列表一样手动索引数据集training_data[index]。我们使用matplotlib来可视化训练数据中的一些样本。

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

2.3为文件创建自定义数据集

自定义的 Dataset 类必须实现三个函数: __init__, __len__, and __getitem__。让我们看看这个实现；FashionMNIST图像存储在img_dir目录中，而它们的标签则单独存储在CSV文件annotations_file中。

import os
import pandas as pd
from torchvision.io import read_image

class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

接下来，将分解每个函数中发生的事情：

__init__函数。类的构造函数，在实例化 Dataset 对象时运行。
__len__函数。返回数据集中的样本数。
__getitem__函数。从数据集中的给定索引idx加载并返回一个示例。根据索引，它识别图像在磁盘上的位置，使用read_image将其转换为张量，从self.img_labels中的csv数据中检索相应的标签，调用它们上的转换函数（如果适用），并以元组形式返回张量图像和相应的标签。

2.4使用DataLoader为训练准备数据

Dataset 检索数据集的特征，并一次标记一个样本。在训练模型时，通常希望以“小批处理minibatches”的形式传递样本，在每个epoch重新洗牌数据以减少模型过度拟合，并使用Python的 multiprocessing 来加快数据检索。

DataLoader是一个 iterable，它用一个简单的API为我们抽象了这种复杂性。

from torch.utils.data import DataLoader

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

2.5遍历数据加载器

上面已将该数据集加载到 DataLoader中，并可根据需要遍历该数据集。下面的每次迭代都会返回一批 train_features 和 train_labels(分别包含 batch_size=64 个特征和标签)。由于指定了 shuffle=True，因此在遍历完所有批次后，数据将被重新洗牌(如需对数据加载顺序进行更精细的控制,请查看 Samplers )。

# Display image and label.
# 从train_dataloader中获取一个批次的图像数据和对应的标签
train_features, train_labels = next(iter(train_dataloader))

# 打印图像数据和标签的形状
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")

# 获取第一个图像，并去除可能存在的批次维度
img = train_features[0].squeeze()

# 获取第一个标签
label = train_labels[0]

# 显示图像
plt.imshow(img, cmap="gray")
plt.show()

# 打印标签
print(f"Label: {label}")

3.Transforms

数据并不总是以训练机器学习算法所需的最终处理形式出现。这时使用transforms来执行数据的一些操作，并使其适合于训练。所有 TorchVision 数据集都有两个参数：transform 用于修改特征， target_transform 用于修改标签。它们接受包含转换逻辑的可调用对象。torchvision.transforms 模块提供了几种常用的转换，可以直接使用。

FashionMNIST 的特征以 PIL 图像格式提供，而标签则为整数。在训练过程中，需要将特征转换为标准化的张量，并将标签转换为 one-hot 编码的张量。为了进行这些转换，使用 ToTensor 和 Lambda。

import torch
from torchvision import datasets
from torchvision.transforms import ToTensor, Lambda

ds = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
    target_transform=Lambda(lambda y: torch.zeros(10, dtype=torch.float).scatter_(0, torch.tensor(y), value=1))
)

ToTensor()函数：ToTensor 将 PIL 图像或 NumPy ndarray 转换为 FloatTensor，并将图像的像素强度值缩放到范围 [0., 1.] 内。

Lambda Transforms：Lambda transforms 应用用户定义的 lambda 函数。在这里，定义了一个函数将整数转换为 one-hot 编码的张量。它首先创建一个大小为 10 的零张量（数据集中标签的数量），然后调用 scatter_ 函数，该函数将值为 1 分配到由标签 y 给出的索引上。

4.构建一个神经网络

神经网络由执行数据操作的层/模块组成。 torch.nn 命名空间提供了构建自己的神经网络所需的所有构建模块。PyTorch 中的每个模块都是 nn.Module 的子类。神经网络本身也是一个模块，它由其他模块（层）组成。

接下来将构建一个神经网络来对FashionMNIST数据集中的图像进行分类。

4.1导入包

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

4.2检查GPU是否可用

希望能够在可用的硬件加速器上（如GPU或MPS）上训练我们的模型。接下来检查torch.cuda or torch.backends.mps 是否可用。

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

4.3定义类

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

import torch.nn as nn

class NeuralNetwork(nn.Module):
    def __init__(self):
        # 初始化神经网络结构
        super().__init__()  # 调用父类的构造函数
        self.flatten = nn.Flatten()  # 将输入图像展平为一维向量
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),  # 输入大小为28*28，输出大小为512的全连接层
            nn.ReLU(),  # ReLU激活函数
            nn.Linear(512, 512),  # 输入和输出大小均为512的全连接层
            nn.ReLU(),  # ReLU激活函数
            nn.Linear(512, 10),  # 输入大小为512，输出大小为10的全连接层（用于10个类别的分类任务）
        )

    def forward(self, x):
        # 定义前向传播过程
        x = self.flatten(x)  # 将输入图像展平
        logits = self.linear_relu_stack(x)  # 经过线性层和ReLU激活函数的堆叠
        return logits  # 返回最终输出

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

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

要使用模型，需要将输入数据传递给它。这将执行模型的 forward 方法，以及一些 background operations 。不要直接调用 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) # 表明Softmax操作应该沿着张量logits的第2个维度进行
y_pred = pred_probab.argmax(1) # argmax(1)表示沿着张量的第2个维度找到最大值所在的索引，这样就可以确定每个样本预测的类别。
print(f"Predicted class: {y_pred}")

4.4模型层

接下来，让我们逐层分解 FashionMNIST 模型中的层。为了说明这一点，将取一个大小为 28x28 的样本小批量，其中包含 3 张图像，并看看当我们将其通过网络时会发生什么。

input_image = torch.rand(3,28,28)
print(input_image.size())

4.4.1nn.Flatten

我们初始化 nn.Flatten 层，将每个 2D 的 28x28 图像转换为一个包含 784 个像素值的连续数组（小批量维度保持（在 dim=0 处））。

flatten = nn.Flatten()
flat_image = flatten(input_image)
print(flat_image.size())

4.4.2nn.Linear

linear layer 是一个模块，它使用其存储的权重和偏置对输入进行线性变换。

layer1 = nn.Linear(in_features=28*28, out_features=20)
hidden1 = layer1(flat_image)
print(hidden1.size())

4.4.3nn.ReLU

非线性激活函数是创建模型输入和输出之间复杂映射的关键。它们在线性变换之后应用，引入非线性，帮助神经网络学习各种现象。

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

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

4.4.4nn.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.5nn.Softmax

神经网络的最后一个线性层返回的是 logits，即在 [-infty, infty] 范围内的原始值，这些值会被传递到 nn.Softmax 模块。logits 会被缩放到 [0, 1] 范围内的值，表示每个类别的模型预测概率。dim参数指示值必须在其指定的维度上求和为 1。

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

4.5模型参数

神经网络中的许多层都是参数化的，即在训练期间进行优化的相关权重和偏置。通过对 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")

5.使用 torch.autograd 进行自动微分

在训练神经网络时，最常用的算法是反向传播。在这个算法中，参数（模型权重）根据损失函数相对于给定参数的梯度进行调整。要计算这些梯度，PyTorch 有一个内置的微分引擎叫做 torch.autograd。它支持对任何计算图进行梯度的自动计算。

考虑最简单的单层神经网络，其中包含输入 x、参数 w 和 b，以及一些损失函数。可以在 PyTorch 中如下定义它：

import torch

x = torch.ones(5)  # input tensor
y = torch.zeros(3)  # expected output
w = torch.randn(5, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)
z = torch.matmul(x, w)+b
loss = torch.nn.functional.binary_cross_entropy_with_logits(z, y)

5.1张量，函数与计算图

此代码定义以下计算图：

在这个网络中，w 和 b 是需要优化的参数。因此，我们需要能够计算损失函数相对于这些变量的梯度。为了实现这一点，我们将这些张量的 requires_grad 属性设置为 True。(注：可以在创建张量时设置 requires_grad 的值，或者稍后使用 x.requires_grad_(True) 方法进行设置。)

我们应用于张量以构建计算图的函数实际上是 Function 类的对象。该对象知道如何在前向方向计算函数，也知道在反向传播步骤中如何计算它的导数。对于反向传播函数的引用存储在张量的 grad_fn 属性中。

print(f"Gradient function for z = {z.grad_fn}")
print(f"Gradient function for loss = {loss.grad_fn}")

5.2计算梯度

为了优化神经网络中参数的权重，我们需要计算损失函数相对于参数的导数，即在一些固定的 x 和 y 值下，需要计算 ${\frac{\partial loss}{\partial w}}{\mathrm{~and~}}{\frac{\partial loss}{\partial b}}$ 。为了计算这些导数，我们调用 loss.backward()，然后从 w.grad 和 b.grad 中检索值：

loss.backward()
print(w.grad)
print(b.grad)

[!CAUTION]

我们只能获取计算图的叶节点的 grad 属性，这些叶节点的 requires_grad 属性设置为 True。对于计算图中的所有其他节点，梯度将不可用。

出于性能原因，我们只能在给定图上执行一次backward的梯度计算。如果我们需要在同一图上进行多次backward传播调用，则需要在 backward 调用中传递 retain_graph=True。

5.3禁用梯度跟踪

默认情况下，所有 requires_grad=True 的张量都在跟踪它们的计算历史并支持梯度计算。然而，有些情况下我们不需要这样做，例如，当我们已经训练好模型，只想将其应用到一些输入数据时，即我们只想通过网络进行前向计算。可以通过将我们的计算代码包裹在 torch.no_grad() 块中来停止跟踪计算：

z = torch.matmul(x, w)+b
print(z.requires_grad)

with torch.no_grad():
    z = torch.matmul(x, w)+b
print(z.requires_grad)

实现相同结果的另一种方法是使用张量的 detach() 方法：

z = torch.matmul(x, w)+b
z_det = z.detach()
print(z_det.requires_grad)

有一些原因你可能想要禁用梯度跟踪：

将神经网络中的一些参数标记为冻结参数（ frozen parameters.）。
当你只进行前向传递时，加快计算速度，因为不跟踪梯度的张量上的计算会更高效。

5.4计算图更多信息

概念上，autograd 在一个由 Function 对象组成的有向无环图（DAG）中记录了数据（张量）和所有执行的操作（以及生成的新张量）。在这个 DAG 中，叶节点是输入张量，根节点是输出张量。通过从根节点到叶节点追踪这个图，可以使用链式法则自动计算梯度。

在前向传播中，autograd 同时执行两件事情：

运行请求的操作以计算结果张量。
在 DAG 中维护操作的梯度函数。

当在 DAG 根节点上调用 .backward() 时，反向传播开始。然后，autograd：

从每个 .grad_fn 计算梯度。
将它们累积在各自张量的 .grad 属性中。
使用链式法则一直传播到叶节点张量。

在 PyTorch 中，DAGs 是动态的。一个重要的事情要注意的是，图是从头开始重新创建的；在每次 .backward() 调用之后，autograd 开始填充一个新图。这正是允许您在模型中使用控制流语句的原因；如果需要，您可以在每次迭代中更改形状、大小和操作。

5.5张量梯度和雅可比积（可选）

在许多情况下，我们有一个标量损失函数，并且我们需要计算相对于一些参数的梯度。然而，有些情况下输出函数是任意张量。在这种情况下，PyTorch 允许您计算所谓的雅可比积（Jacobian product），而不是实际的梯度。

对于向量函数 $\vec{y}=f(\vec{x})$ ，其中 $\vec{x}=\langle x_1,\dots,x_n\rangle$ 且 $\vec{y}=\langle y_1,\dots,y_m\rangle$ ， $\vec{y}$ 相对于 $\vec{x}$ 的梯度由雅可比矩阵给出：

$J=\left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{1}}{\partial x_{n}} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_{m}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)$

而不是计算雅可比矩阵本身，PyTorch 允许您计算给定输入向量 $\mathbf{v}=(v_1,\dots,v_m)$ 的雅可比积 $\mathbf{v}^T\cdot J$ 。这可以通过使用 $\mathbf{v}$ 作为参数调用 backward 来实现。 $\mathbf{v}$ 的大小应与我们想要计算乘积的原始张量的大小相同：

# 创建了一个大小为4x5的单位矩阵张量 inp，并将其设置为需要梯度信息。
inp = torch.eye(4, 5, requires_grad=True)

# 定义了一个计算图，将输入张量 inp 的每个元素加1，然后对结果取平方，并转置得到输出张量 out。
out = (inp+1).pow(2).t()

# 执行反向传播，计算输出张量 out 对自身的梯度。这里使用了全1的梯度张量作为参数，表示将梯度传播到 out 中的每个元素。retain_graph=True 保留了计算图，以便后续再次执行反向传播。
out.backward(torch.ones_like(out), retain_graph=True)
# 打印第一次反向传播后输入张量 inp 的梯度。由于PyTorch会累积梯度信息，因此这里会显示第一次反向传播后的梯度值。
print(f"First call\n{inp.grad}")

# 再次执行反向传播，计算输出张量 out 对自身的梯度。由于之前已经调用过一次反向传播，所以这里会继续累积梯度信息。
out.backward(torch.ones_like(out), retain_graph=True)
# 打印第二次反向传播后输入张量 inp 的梯度。这里会显示第一次和第二次反向传播后的梯度值的累积结果。
print(f"\nSecond call\n{inp.grad}")

# 将输入张量 inp 的梯度信息清零，以便后续重新累积梯度。
inp.grad.zero_()

# 再次执行反向传播，计算输出张量 out 对自身的梯度。由于之前调用了 inp.grad.zero_()，所以这里的梯度信息会重新累积。
out.backward(torch.ones_like(out), retain_graph=True)
# 打印清零梯度后的结果，显示输入张量 inp 的梯度信息已经被重新累积。
print(f"\nCall after zeroing gradients\n{inp.grad}")

注意，当我们第二次使用相同的参数调用 backward 时，梯度的值是不同的。这是因为在进行 backward 传播时，PyTorch 累积梯度，即计算的梯度值被加到计算图的所有叶节点的 grad 属性中。如果您想要计算正确的梯度，您需要在此之前将 grad 属性清零。在实际训练中，优化器帮助我们完成这一点。

注意：之前我们没有带参数调用 backward() 函数。这本质上等同于调用 backward(torch.tensor(1.0))，这是在标量值函数（例如神经网络训练中的损失）情况下计算梯度的一种有用方式。

6.优化模型参数

既然我们已经有了模型和数据，现在应该训练、验证、测试我们的模型(基于我们的数据来优化参数)。训练一个模型也是一个迭代的过程；在每次迭代中(又称为 epoch)，模型会对输出中进行一次预测，计算这个预测的误差(损失值)，收集这些误差相对于参数的导数，然后通过梯度下降的方式来优化这些参数。关于这个过程的更详细的介绍，可以看3Blue1Brown制作的《反向传播演算》这个视频。

6.1前提代码

将之前的模块（数据集和数据加载器、构建模型）的代码拿过来。

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)

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

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

model = NeuralNetwork()

6.2超参数

超参数是你用来控制模型优化过程的、可以调整的参数。不同的超参数取值能够影响模型训练和收敛的速度(更多关于调整超参数的内容)

定义以下用于训练的超参数:

Number of Epochs - 迭代数据集的次数
Batch Size - 在参数更新之前通过网络传播的数据样本数量。
Learning Rate - 学习率，每 Batch/Epoch 次更新模型参数的幅度。较小的值会产生较慢的学习速度，较大的值可能会在训练过程中产生无法预料的行为。

learning_rate = 1e-3
batch_size = 64
epochs = 5

6.3优化循环

设置完超参数后，接下来我们在一个优化循环中训练并优化我们的模型。优化循环的每次迭代叫做一个 Epoch(时期、纪元)。

每个 Epoch 由两个主要部分构成:

训练循环 在训练数据集上遍历，尝试收敛到最优的参数。
验证/测试循环 在测试数据集上遍历，以检查模型效果是否在提升。

接下来，让我们简单的熟悉一下在训练循环中使用的一些概念。

6.4损失函数

拿到一些训练数据的时候，我们的模型不太可能给出正确答案。损失函数能衡量获得的结果相对于目标值的偏离程度，我们希望在训练中能够最小化这个损失函数。我们对给定的数据样本做出预测然后和真实标签数据对比来计算损失。

常见的损失函数包括给回归任务用的 nn.MSELoss(Mean Square Error, 均方误差)、给分类任务使用的 nn.NLLLoss(Negative Log Likelihood, 负对数似然)、nn.CrossEntropyLoss(交叉熵损失函数)结合了 nn.LogSoftmax 和 nn.NLLLoss.

我们把模型输出的 logits 传递给 nn.CrossEntropyLoss，它会正则化 logits 并计算预测误差。

# 初始化损失函数
loss_fn = nn.CrossEntropyLoss()

6.5优化器

优化是在每一个训练步骤中调整模型参数来减小模型误差的过程。优化算法定义了这个过程应该如何进行(在这个例子中，我们使用 Stochastic Gradient Descent-即SGD，随机梯度下降)。所有优化的逻辑都被封装在 optimizer 这个对象中。这里，我们使用 SGD 优化器。除此之外，在 PyTorch 中还有很多其他可用的优化器，比如 ADAM 和 RMSProp 在不同类型的模型和数据上表现得更好。

通过注册需要训练的模型参数、然后传递学习率这个超参数来初始化优化器。

optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

在训练循环内部, 优化在三个步骤上发生：

调用 optimizer.zero_grad() 来重置模型参数的梯度。梯度会默认累加，为了防止重复计算(梯度)，我们在每次迭代中显式的清空(梯度累加值)。
调用 loss.backward() 来反向传播预测误差。PyTorch 对每个参数分别存储损失梯度。
获取到梯度后，调用 optimizer.step() 来根据反向传播中收集的梯度来调整参数。

6.6完整实现

定义 train_loop 为优化循环的代码，test_loop 为根据测试数据来评估模型表现的代码。

def train_loop(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)
    # 将模型设置为训练模式-对于 batch normalization 和 dropout 层很重要
    # 在这种情况下不需要，但为最佳实践添加了
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        # Compute prediction and loss
        pred = model(X)  # 计算模型的预测值
        loss = loss_fn(pred, y)  # 计算损失

        # Backpropagation
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 根据梯度更新模型参数
        optimizer.zero_grad()  # 清除梯度

        if batch % 100 == 0:
            # 打印损失信息
            loss, current = loss.item(), (batch + 1) * len(X)
            print(f"loss: {loss:>7f}  [{current:>5d}/{size:>5d}]")


def test_loop(dataloader, model, loss_fn):
    # 将模型设置为评估模式-对于 batch normalization 和 dropout 层很重要
    # 在这种情况下不需要，但为最佳实践添加了
    model.eval()
    size = len(dataloader.dataset)
    num_batches = len(dataloader) 
    test_loss, correct = 0, 0

    # 使用 torch.no_grad() 对模型进行评估，确保在测试模式下不计算梯度
    # 同时也减少了对 requires_grad=True 的张量进行不必要的梯度计算和内存使用
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()  # 计算测试集上的损失
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()  # 统计正确预测的数量

    test_loss /= num_batches  # 计算平均损失
    correct /= size  # 计算准确率
    print(f"Test Error: \n Accuracy: {(100*correct):>0.1f}%, Avg loss: {test_loss:>8f} \n")

初始化损失函数和优化器，传递给 train_loop 和 test_loop。可以随意地修改 epochs 的数量来跟踪模型表现的进步情况。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_loop(train_dataloader, model, loss_fn, optimizer)
    test_loop(test_dataloader, model, loss_fn)
print("Done!")

7.模型保存和加载

在这个章节我们会学习如何持久化模型状态来保存、加载和执行模型预测。

import torch
import torchvision.models as models

7.1模型权重的保存和加载

PyTorch 将模型学习到的参数存储在一个内部状态字典中，叫 state_dict。它们可以通过 torch.save 方法来持久化。

# 使用 models.vgg16(weights='IMAGENET1K_V1') 加载一个预训练的 VGG16 模型，
# 该模型使用 ImageNet 数据集进行了训练。
model = models.vgg16(weights='IMAGENET1K_V1')

# 调用 torch.save() 函数来保存这个模型的参数到一个文件中，
# 文件名为 'model_weights.pth'。
torch.save(model.state_dict(), 'model_weights.pth')

要加载模型权重，需要先创建一个跟要加载权重的模型结构一样的模型，然后使用 load_state_dict() 方法加载参数。

# 创建一个未经训练的 VGG16 模型，不指定 ``weights``
model = models.vgg16()

# 加载之前保存的模型参数到模型中
model.load_state_dict(torch.load('model_weights.pth'))

# 将模型设置为评估模式
model.eval()