PyTorch官方文档-60分钟入门

这是官方入门文档，由于这些框架更新很快，建议先看官方文档，上面应该最新的代码写法。

同时建议看B站的一个 PyTorch 视频教程（网页二），短小精悍，用于实操不太行，但是有利于深入理解 PyTorch 代码原理。

张量

PyTorch 输入输出，参数等等均是 Tensor 数据结构，类似与 Numpy ，除了 Tensors 可以在 GPU 上进行加速这一点不同外。

首先加载包

import torch
import numpy as np

初始化张量

有多种方式构建一个张量

• 从数据直接构建张量，并且可以自动推断数据类型

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

• 从 NumPy 数组构建张量

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

• 从另一个张量中构建新张量

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

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

• 构建随机张量或固定值的张量

  这里 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}")

张量属性

属性包括 shape ，数据类型，存储在什么设备上。

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

张量操作

张量有超过100个操作方法，包括算数、线性代数、矩阵操作（转置、索引、切片）、采样等

每个操作都可以在GPU上运行（通常比在CPU上速度更快)。如果你在使用Colab，可以通过修改 Edit > Notebook Settings 来分配一个 GPU。

# We move our tensor to the GPU if available
device = torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu'
tensor = tensor.to(device)
print(f"Device tensor is stored on: {tensor.device}")

尝试下列操作，如果你已经对NumPy API十分熟悉，上手张量API将会很简单。

类似numpy的索引和切片操作

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

连接张量 你可以使用 torch.cat 沿着给定的维度连接一系列张量。另一个张量连接操作符， 与 torch.cat 稍有不同，请参阅 torch.stack。

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

运算操作

# This computes the matrix multiplication between two tensors. y1, y2, y3 will have the same value
# ``tensor.T`` returns the transpose of a tensor
y1 = tensor @ tensor.T
y2 = tensor.matmul(tensor.T)

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


# This computes the element-wise product. z1, z2, z3 will have the same value
z1 = tensor * tensor
z2 = tensor.mul(tensor)

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

单个元素的张量 在聚合运算场景中，你可能会得到一个单元素的张量，可使用item()将其传唤为Python数值。

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

原地操作 修改张量中的原值操作称为原地操作。它们以 _ 后缀表示。例如：x.copy_(y)，x.t_()，会改变 x。

print(f"{tensor} \n")
tensor.add_(5)
print(tensor)

与NumPy转换

CPU 上的张量和 NumPy 数组可以共享它们的底层内存位置，改变其中一个就会改变另一个。

张量转为NumPy array

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

对于张量的修改体现到了NumPy array上。

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

NumPy array转为张量

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

NumPy array转为张量（对于 NumPy array 的修改也体现到了对张量的修改）

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

torch.autograd 的简明介绍

torch.autograd 是 PyTorch 的自动微分引擎，用于神经网络的训练。在本节中，您将了解 autograd 如何帮助神经网络进行训练。

背景

神经网络（NNs）是一组嵌套函数，对某些输入数据进行处理。这些函数由参数（包括权重和偏置）定义，在PyTorch中以张量的形式存储。

训练神经网络分为两个步骤：

• 前向传播：在前向传播中，神经网络会根据输入数据对正确输出做出最佳猜测。它通过依次运行各个函数来实现这一猜测。
• 反向传播：在反向传播中，神经网络会根据其预测误差调整参数。具体做法是从输出层向后逐层反向计算，收集误差对各函数参数的导数（梯度），然后通过梯度下降法优化这些参数。

在 PyTorch 中的使用

我们来看一个训练步骤。在这个例子中，我们从 torchvision 加载一个预训练的 resnet18 模型。我们创建一个随机数据张量，表示一张包含 3 个通道、高和宽均为 64 的图像，并将对应的标签初始化为一些随机值。预训练模型中的标签形状为 (1,1000)。

注意：本教程仅适用于CPU，不适用于GPU设备（即使张量被移动到CUDA）。

import torch
from torchvision.models import resnet18, ResNet18_Weights
model = resnet18(weights=ResNet18_Weights.DEFAULT)
data = torch.rand(1, 3, 64, 64)
labels = torch.rand(1, 1000)

这里图像数据维度为 形状为 (N, C, H, W) ，这里 N 是批次大小。

接下来，我们将输入数据通过模型的每一层进行处理，以生成预测结果。这称为前向传播。

prediction = model(data) # forward pass

我们使用模型的预测结果和对应的标签来计算误差（loss）。下一步是将该误差通过网络进行反向传播。当我们在误差张量上调用 .backward() 时，反向传播即开始执行。随后，Autograd 会计算并存储每个模型参数的梯度，这些梯度被保存在参数的 .grad 属性中。

loss = (prediction - labels).sum()
loss.backward() # backward pass

接下来，我们加载一个优化器，在本例中使用学习率为0.01、动量为0.9的随机梯度下降（SGD）。我们将模型的所有参数注册到该优化器中。

optim = torch.optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

最后，我们调用 .step() 来启动梯度下降。优化器会根据 .grad 中存储的梯度来调整每个参数。

optim.step() #gradient descent

此时，你已经具备了训练神经网络所需的一切。以下部分详细介绍了自动求导（autograd）的工作原理，你可以选择跳过。

自动求导中的微分

让我们来看看 autograd 如何收集梯度。我们创建了两个具有 requires_grad=True 的张量 a 和 b，这向 autograd 发出信号，告知它对这两个张量的每一项操作都应被跟踪。

import torch

a = torch.tensor([2., 3.], requires_grad=True)
b = torch.tensor([6., 4.], requires_grad=True)

我们从 a 和 b 构造另一个张量 Q。 \[ Q = 3 a^3 - b^2 \]

Q = 3*a**3 - b**2

假设 a 和 b 是神经网络的参数，Q 是误差。在神经网络训练中，我们希望得到误差相对于参数的梯度，即 \[ \frac{\partial Q}{\partial a} = 9a^2 \\ \frac{\partial Q}{\partial b} = -2b \] 当我们对 Q 调用 .backward() 时，自动求导会计算这些梯度，并将它们存储在相应张量的 .grad 属性中。

因为 Q 是一个向量，因此我们需要在 Q.backward() 中显式传递 gradient 参数（因为一般计算梯度是针对标量损失的，针对向量传入的 gradient 参数相当于告诉 PyTorch：假设你有一个标量损失 \(L\)，且 \(\frac{∂L}{∂\boldsymbol{y}}=\boldsymbol{v}\)，这里 \(\boldsymbol{v}\) 就是提供的参数，请帮我计算 \(\frac{∂L}{∂\boldsymbol{x}}\) ）。

gradient 是一个与 Q 具有相同形状的张量，表示对 Q 进行求导的结果，如 \[ \frac{dQ}{dQ} = 1 \]

因此，这里提供的 gradient 参数是一个全为1的与Q 具有相同形状的张量。

external_grad = torch.tensor([1., 1.])
Q.backward(gradient=external_grad)

同样地，我们也可以将 Q 集成一个标量，并隐式调用反向传播，例如 Q.sum().backward()。

梯度现在被存储在 a.grad 和 b.grad 中

# check if collected gradients are correct
print(9*a**2 == a.grad)
print(-2*b == b.grad)

tensor([True, True])
tensor([True, True])

选读部分：使用 autograd 进行向量微分

数学上，若有向量函数 \(\boldsymbol{y} = \boldsymbol{f(x)}\) ，那么 \(\boldsymbol{y}\) 关于 \(\boldsymbol{x}\) 的梯度就是一个雅可比矩阵：
\[ \boldsymbol{J} = \left( \frac{\partial \boldsymbol{y}}{\partial x_1} \quad \dots \quad \frac{\partial \boldsymbol{y}}{\partial x_n} \right) = \begin{pmatrix} \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{pmatrix} \] 而 torch.autograd 这个包就是用来计算一些向量与雅可比矩阵的乘积（\(\boldsymbol{v} \boldsymbol{J}\) ）的。例如，如果 \(\boldsymbol{v}\) 是一个标量函数 \(l = g(\boldsymbol{y})\) 的梯度：\(\boldsymbol{v} = \begin{pmatrix} \dfrac{\partial l}{\partial y_1} & \cdots & \dfrac{\partial l}{\partial y_m} \end{pmatrix}\) 由链式法则，我们可以得到： \[ \boldsymbol{v} \boldsymbol{J} = \begin{pmatrix} \dfrac{\partial l}{\partial y_1} & \cdots & \dfrac{\partial l}{\partial y_m} \end{pmatrix} \begin{pmatrix} \dfrac{\partial y_1}{\partial x_1} & \cdots & \dfrac{\partial y_1}{\partial x_n} \\ \vdots & \ddots & \vdots \\ \dfrac{\partial y_m}{\partial x_1} & \cdots & \dfrac{\partial y_m}{\partial x_n} \end{pmatrix} = \begin{pmatrix} \dfrac{\partial l}{\partial x_1} & \cdots & \dfrac{\partial l}{\partial x_n} \end{pmatrix} \] 官网例子中 \(\boldsymbol{v}\) 设置为竖向量，因此这里改成了 \(\boldsymbol{J}^{T} \boldsymbol{v}\) ，不过是一个意思。

上面的 gradient 参数就是这里的 \(\boldsymbol{v}\) 向量。

注意：grad在反向传播过程中是累加的(accumulated)，这意味着每一次运行反向传播，梯度都会累加之前的梯度，所以一般在反向传播之前需把梯度清零。

计算图

从概念上讲，自动求导通过由函数对象构成的有向无环图（DAG），记录数据（张量）以及所有执行的操作（包括生成的新张量）。在这个DAG中，叶子节点是输入张量，根节点是输出张量。通过从根节点追踪到叶子节点，可以利用链式法则自动计算梯度。

在前向传播过程中，自动求导会同时完成两件事：

• 执行所请求的操作以计算结果张量，并
• 在DAG中保持操作的梯度函数。

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

• 计算每个 .grad_fn 的梯度
• 将它们累积到相应张量的 .grad 属性中，并
• 使用链式法则，一路传播到叶节点张量。

以下是本例中DAG的可视化表示。图中的箭头方向指向前向传播的方向。节点代表前向传播中每个操作的反向函数。蓝色的叶节点表示我们的叶张量a和b。

在 PyTorch 中，DAG 是动态的。需要注意的是，图会从头开始重新构建；每次调用 .backward() 之后，自动求导（autograd）就会开始填充一个新的图。这正是允许你在模型中使用控制流语句的原因：你可以根据需要，在每一轮迭代中修改图的形状、大小和操作。

排除在DAG之外的张量

torch.autograd 会跟踪所有 requires_grad 标志设为 True 的张量上的操作。对于不需要梯度的张量，将该属性设置为 False 会将其排除在梯度计算 DAG 之外。

即使只有一个输入张量的 requires_grad=True，输出张量仍需要梯度。

x = torch.rand(5, 5)
y = torch.rand(5, 5)
z = torch.rand((5, 5), requires_grad=True)

a = x + y
print(f"Does `a` require gradients?: {a.requires_grad}")
b = x + z
print(f"Does `b` require gradients?: {b.requires_grad}")

Does `a` require gradients?: False
Does `b` require gradients?: True

在神经网络中，那些不计算梯度的参数通常被称为冻结参数。如果你事先知道某些参数不需要梯度，那么将这些参数“冻结”是有用的（这能带来一些性能优势，通过减少自动求导计算来实现）。

在微调过程中，我们冻结模型的大部分参数，通常只修改分类器层以对新标签进行预测。下面我们通过一个小例子来演示这一过程。与之前一样，我们加载一个预训练好的ResNet18模型，并冻结所有参数。

from torch import nn, optim

model = resnet18(weights=ResNet18_Weights.DEFAULT)

# Freeze all the parameters in the network
for param in model.parameters():
    param.requires_grad = False

假设我们想在包含10个标签的新数据集上对模型进行微调。在ResNet中，分类器是最后一个线性层 model.fc。我们可以直接用一个新的线性层（默认不冻结）来替换它，该层将作为我们的分类器。

model.fc = nn.Linear(512, 10)

现在模型中的所有参数（除了 model.fc 的参数）都被冻结了。唯一能计算梯度的参数是 model.fc 的权重和偏置。

# Optimize only the classifier
optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)

请注意，尽管我们在优化器中注册了所有参数，但只有分类器的权重和偏置会计算梯度（因此在梯度下降中被更新）。

神经网络

可以使用 torch.nn 包来构建神经网络。

既然你已经了解了自动求导，神经网络（nn）就依赖自动求导来定义模型并对其进行求导。一个 nn.Module 包含多个层，以及一个名为 forward(input) 的方法，用于返回输出结果 output。

例如，看看这个用于分类数字图像的网络：

这是一个简单的前馈网络。它接收输入，依次通过多个层，最后输出结果。

神经网络的典型训练过程如下：

• 定义具有一些可学习参数 (或权重) 的神经网络

• 迭代输入数据集

• 通过网络处理输入

• 计算损失 (输出与正确值的差距)

• 将梯度反向传播回网络的参数

• 更新网络的权重，通常使用一个简单的更新规则: weight = weight - learning_rate * gradient

定义网络

定义网络如下

import torch
import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):

    def __init__(self):
        super().__init__()
        # 1 input image channel, 6 output channels, 5x5 square convolution
        # kernel
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 5*5 from image dimension
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, input):
        # Convolution layer C1: 1 input image channel, 6 output channels,
        # 5x5 square convolution, it uses RELU activation function, and
        # outputs a Tensor with size (N, 6, 28, 28), where N is the size of the batch
        c1 = F.relu(self.conv1(input))
        # Subsampling layer S2: 2x2 grid, purely functional,
        # this layer does not have any parameter, and outputs a (N, 6, 14, 14) Tensor
        s2 = F.max_pool2d(c1, (2, 2))
        # Convolution layer C3: 6 input channels, 16 output channels,
        # 5x5 square convolution, it uses RELU activation function, and
        # outputs a (N, 16, 10, 10) Tensor
        c3 = F.relu(self.conv2(s2))
        # Subsampling layer S4: 2x2 grid, purely functional,
        # this layer does not have any parameter, and outputs a (N, 16, 5, 5) Tensor
        s4 = F.max_pool2d(c3, 2)
        # Flatten operation: purely functional, outputs a (N, 400) Tensor
        s4 = torch.flatten(s4, 1)
        # Fully connected layer F5: (N, 400) Tensor input,
        # and outputs a (N, 120) Tensor, it uses RELU activation function
        f5 = F.relu(self.fc1(s4))
        # Fully connected layer F6: (N, 120) Tensor input,
        # and outputs a (N, 84) Tensor, it uses RELU activation function
        f6 = F.relu(self.fc2(f5))
        # Fully connected layer OUTPUT: (N, 84) Tensor input, and
        # outputs a (N, 10) Tensor
        output = self.fc3(f6)
        return output


net = Net()
print(net)

输出结果为

Net(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)

你只需定义前向函数，后向函数（用于计算梯度）将通过自动微分（autograd）自动生成。在前向函数中可以使用任意的张量操作。

模型的可学习参数通过 net.parameters() 返回

params = list(net.parameters())
print(len(params))
print(params[0].size())  # conv1's .weight

10
torch.Size([6, 1, 5, 5])

我们来尝试一个随机的32x32输入。注意：该网络（LeNet）的预期输入尺寸为32x32。若要使用此网络处理MNIST数据集，请将数据集中的图像调整为32x32大小。

input = torch.randn(1, 1, 32, 32)
out = net(input)
print(out)

tensor([[-0.0371,  0.0160, -0.1203, -0.0563, -0.1361, -0.0697, -0.0777, -0.0792,
         -0.1307,  0.0206]], grad_fn=<AddmmBackward0>)

将所有参数的梯度缓冲区清零，并使用随机梯度进行反向传播：

net.zero_grad()
out.backward(torch.randn(1, 10))

注意： torch.nn 仅支持小批量（mini-batches）输入。整个 torch.nn 包只支持以小批量样本的输入，不支持单个样本。

例如，nn.Conv2d 会接收一个形状为 nSamples × nChannels × Height × Width 的四维张量。

如果你只有一个样本，只需使用 input.unsqueeze(0) 来添加一个虚拟的批次维度。

损失函数

损失函数接收输入的（输出，目标）对，并计算一个值，用于估计输出与目标之间的距离。

nn 包中包含多种不同的损失函数。一个简单的损失函数是：nn.MSELoss，它计算输出与目标之间的均方误差。

举个例子

output = net(input)
target = torch.randn(10)  # a dummy target, for example
target = target.view(1, -1)  # make it the same shape as output
criterion = nn.MSELoss()

loss = criterion(output, target)
print(loss)

tensor(2.4473, grad_fn=<MseLossBackward0>)

现在，如果你沿着反向方向跟踪损失，并使用其 .grad_fn 属性，你将看到一个类似这样的计算图：

input -> conv2d -> relu -> maxpool2d -> conv2d -> relu -> maxpool2d
      -> flatten -> linear -> relu -> linear -> relu -> linear
      -> MSELoss
      -> loss

因此，当我们调用 loss.backward() 时，整个图会相对于神经网络的参数进行求导，图中所有 requires_grad=True 的张量都会将它们的 .grad 张量与梯度相加。

为了说明，让我们倒退几步（没看懂）：

print(loss.grad_fn)  # MSELoss
print(loss.grad_fn.next_functions[0][0])  # Linear
print(loss.grad_fn.next_functions[0][0].next_functions[0][0])  # ReLU

<MseLossBackward0 object at 0x2b48a0d61450>
<AddmmBackward0 object at 0x2b48a0d63880>
<AccumulateGrad object at 0x2b48a0d61450>

反向传播

要反向传播误差，我们只需调用 loss.backward() 即可。不过需要先清除现有的梯度，否则梯度会累加到已有的梯度上。

现在我们将调用 loss.backward()，并查看 conv1 的偏置梯度在反向传播前后的变化。由于我们尚未引入优化器，因此直接在模型上清空梯度。一旦使用了优化器，建议使用 optimizer.zero_grad()，如下所示。

net.zero_grad()     # zeroes the gradient buffers of all parameters

print('conv1.bias.grad before backward')
print(net.conv1.bias.grad)

loss.backward()

print('conv1.bias.grad after backward')
print(net.conv1.bias.grad)

conv1.bias.grad before backward
None
conv1.bias.grad after backward
tensor([ 0.0410,  0.0061, -0.0045,  0.0083, -0.0020,  0.0066])

现在，我们已经了解了如何使用损失函数。

更新权重

实践中最常用的更新规则是随机梯度下降（SGD）：

weight = weight - learning_rate * gradient

我们可以用简单的 Python 代码来实现这一点：

learning_rate = 0.01
for f in net.parameters():
    with torch.no_grad():
        f -= f.grad * learning_rate

然而，当你使用神经网络时，会希望采用各种不同的优化规则，例如随机梯度下降（SGD）、Nesterov-SGD、Adam、RMSProp等。为此，我们开发了一个小型包：torch.optim，它实现了所有这些方法。使用起来非常简单：

import torch.optim as optim

# create your optimizer
optimizer = optim.SGD(net.parameters(), lr=0.01)

# in your training loop:
optimizer.zero_grad()   # zero the gradient buffers
output = net(input)
loss = criterion(output, target)
loss.backward()
optimizer.step()    # Does the update

训练循环内的五个关键步骤 (Training Loop)

第一步：optimizer.zero_grad() —— 梯度清零

• 作用：将模型中所有参数的梯度缓冲区重置为 0。
• 为什么必须做：PyTorch 默认会累加梯度（为了支持大批量训练或梯度累积）。如果不在每次迭代前清零，上一次计算的梯度会与本次的梯度相加，导致参数更新方向错误。

第二步：output = net(input) —— 前向传播

• 将输入数据 input 喂入神经网络，经过各层计算，得到预测输出 output。

第三步：loss = criterion(output, target) —— 计算损失

• criterion 是预先定义好的损失函数（如交叉熵 CrossEntropyLoss 或均方误差 MSELoss）。
• 比较预测值 output 和真实标签 target，计算出一个标量损失值 loss，衡量模型当前预测的离谱程度。

第四步：loss.backward() —— 反向传播

• 这是深度学习的核心数学环节。PyTorch 自动计算 loss 对计算图中所有叶子节点（即 net.parameters()）的梯度（偏导数）。
• 计算出的梯度会累加到各参数的 .grad 属性中（这就是为什么第一步必须先 zero_grad）。

第五步：optimizer.step() —— 参数更新

• 优化器根据刚刚计算出的 .grad 值，结合所选的算法（SGD），执行参数更新。
• 至此，模型完成了一次迭代学习。

训练一个分类器

这就是全部内容。你已经了解了如何定义神经网络、计算损失以及更新网络的权重。

现在你可能在想，数据呢？

通常，当你需要处理图像、文本、音频或视频数据时，可以使用标准的 Python 包将数据加载到 NumPy 数组中，然后将该数组转换为 PyTorch 的 Tensor。

• 对于图像处理，Pillow、OpenCV 等包非常有用。
• 对于音频，可以使用 scipy 和 librosa 等包
• 对于文本处理，无论是基于原始 Python 或 Cython 的加载方式，还是使用 NLTK 和 SpaCy 都很有帮助。

针对视觉任务，我们专门创建了一个名为 torchvision 的包，其中包含适用于 ImageNet、CIFAR10、MNIST 等常见数据集的数据加载器，以及图像处理的数据转换器，例如 torchvision.datasets 和 torch.utils.data.DataLoader。

这大大提高了便利性，并避免了编写冗余代码。

在本教程中，我们将使用 CIFAR10 数据集。该数据集包含以下类别：‘飞机’、‘汽车’、‘鸟’、‘猫’、‘鹿’、‘狗’、‘青蛙’、‘马’、‘船’和‘卡车’。CIFAR-10 中的图像尺寸为 3×32×32，即每张图像有 3 个通道，分辨率为 32×32 像素的彩色图像。

训练一个图像分类器

我们将按以下步骤进行：

1. 使用 torchvision 加载并归一化 CIFAR10 的训练和测试数据集
2. 定义卷积神经网络
3. 定义损失函数
4. 用训练数据对网络进行训练
5. 在测试数据上测试网络

1 加载并归一化 CIFAR10 数据

如果将所有数据存入内存，然后通过 for 循环遍历，就很可能出现内存不足的情况。

实际做法是分批次处理数据（batches），除此之外可能还需要对数据进行处理，转化为 pytorch 需要的格式。

pytorch 处理数据的三大工具为 ：

• transforms ： 加载数据时的操作，负责转为张量，进行归一化处理等操作。
• Dataset ：包装数据。这是“懒加载”的——除非通过索引（trainset[idx]）请求，它才会真正去硬盘读取第 0 张图片，应用 transform，并返回一个元组 (transformed_image_tensor, label)。否则它不会将数据加载到内存中。这种“用的时候再读”的策略，避免了将整个数据集一次性加载到内存中造成溢出。
• DataLoader ：批量数据迭代器，为模型训练不断提供数据。你可以把 DataLoader 想象成一个自动化的流水线工人：
  • 它按照 batch_size（批量大小）把零散的图片打包成箱。
  • 它根据 shuffle 指令打乱图片顺序，防止模型偷看规律。
  • 它派出 num_workers 个助手（子进程）提前备货，确保 GPU 永远不会饿着肚子等数据。当 GPU 算完当前批次时，下一批数据已经“等在门口”了，从而实现 “读取”和“计算”的流水线并行，大幅提升训练效率。

这里的重点就是按照批次从磁盘中读取需要的数据，而不是将所有数据加载到内存，而是每次都只读取需要的数据。

使用 torchvision，加载 CIFAR10 极为简单。

import torch
import torchvision
from torchvision.transforms import v2

torchvision 数据集的输出是范围在 [0， 1] 的 PILImage 图像。我们将它们转换为范围在 [-1， 1] 的归一化张量（利用 v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))）。

注意：如果您在 Windows 或 MacOS 上运行本教程，并遇到与多处理相关的 BrokenPipeError 或 RuntimeError，请尝试将 torch.utils.data.DataLoader() 的 num_worker 设置为 0。

这里下载上百MB数据，速度很慢，我就没有跑了

这里torchvision.transforms.v2 是 PyTorch 官方推出的新一代数据增强与转换模块，它的功能更强大、速度更快，并且正在逐步取代旧版的 torchvision.transforms。核心 API 变化为，由使用 transforms.ToTensor() 一步到位，拆分为 v2.ToImage() + v2.ToDtype() 两步。

transform = v2.Compose([
    v2.ToImage(),
    v2.ToDtype(torch.float32, scale=True),
    v2.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

batch_size = 4

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

这里要注意的一点是， PyTorch 的分类标签必须从 0 开始，这是为了与 CrossEntropyLoss 函数的内部实现机制兼容。

训练集必须打乱顺序 shuffle=True ，目的有2个：

1. 使得不同 epoch 看到的数据不同
2. 打破“批次偏差”（根本原因）。如果数据未打乱，同一批次内包含高度相似的样本（例如都是猫），那么梯度下降方向是偏斜的，导致收敛不稳定。

让我们展示一些训练图片，只是为了好玩。

import matplotlib.pyplot as plt
import numpy as np

# functions to show an image


def imshow(img):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# get some random training images
dataiter = iter(trainloader)
images, labels = next(dataiter)

# show images
imshow(torchvision.utils.make_grid(images))
# print labels
print(' '.join(f'{classes[labels[j]]:5s}' for j in range(batch_size)))

horse horse ship  truck

I/O瓶颈

下面内容来自 deepseek，看样子需要 SSD。

“频繁从磁盘读取”确实是深度学习训练中一个著名的性能瓶颈，叫做 I/O 瓶颈。对于大规模数据集（比如 ImageNet 或视频数据），SSD 几乎是必需品。SSD 的随机读取速度比机械硬盘快得多，提升10倍以上。

当你遇到以下情况时，请务必考虑 SSD，否则 GPU 利用率会掉到 10% 以下：

1. 数据集总大小 > 内存大小：比如你有 512GB 的图片，但电脑只有 32GB 内存，操作系统缓存不了全部数据，每个 Epoch 都必须重新读磁盘。这时候 SSD 的优势就体现出来了。
2. 大量小图片文件（如几千张 4K 高清图）：随机读取极耗性能，SSD 的 4K 随机读写能力远超 HDD。
3. 数据预处理极其复杂：虽然 transform 是在 CPU 上计算的，但读磁盘和解压（如读取 .jpg、.png 文件）本身消耗 CPU。如果 CPU 因为等磁盘而空闲，说明 I/O 太慢了。（没太看懂，但是这一点好像可以通过增加 num_workers 子进程数目来解决，也就是多用几个 CPU 核心）。

2 定义卷积神经网络

复制神经网络部分中的神经网络，并将其修改为处理3通道图像（而不是之前定义的1通道图像）。

简单来说，必须在 init() 中定义的是“有可训练参数”或“有需要维护的状态”的层；而“纯数学运算”则可以直接在 forward() 里调用。

深度学习训练中，权重需要在内存中持续存在，不能像普通局部变量那样用完就销毁。

• 写进 __init__：层的权重被绑定在 self（即整个模型对象）上。只要这个模型对象还在，这些权重就永久存活在内存中，随叫随到。
• 不写进 __init__：如果你在 forward 里写 conv = nn.Conv2d(...)，这个 conv 在 forward 函数执行完就被垃圾回收（销毁）了。下一次调用 forward 时，又会新建一个完全陌生、权重随机的层，之前算的梯度全白费了。

下面 RELU 等激活函数和 pool 层、flatten 层等没有参数，因此可以不用写入 __init__ 。

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()

3 定义损失函数和优化器

我们使用分类交叉熵损失和动量随机梯度下降。

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

4 训练网络

这时事情就开始变得有趣了。我们只需遍历数据迭代器，将输入传递给网络并进行优化。

训练步骤包括以下部分

1. 梯度清零
2. 前向传播
3. 计算损失
4. 反向传播
5. 参数更新
6. 记录损失并打印出来

for epoch in range(2):  # loop over the dataset multiple times

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if i % 2000 == 1999:    # print every 2000 mini-batches
            print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
            running_loss = 0.0

print('Finished Training')

[1,  2000] loss: 2.210
[1,  4000] loss: 1.927
[1,  6000] loss: 1.744
[1,  8000] loss: 1.605
[1, 10000] loss: 1.531
[1, 12000] loss: 1.474
[2,  2000] loss: 1.397
[2,  4000] loss: 1.384
[2,  6000] loss: 1.337
[2,  8000] loss: 1.307
[2, 10000] loss: 1.307
[2, 12000] loss: 1.278
Finished Training

让我们快速保存训练好的模型：

PATH = './cifar_net.pt'
torch.save(net.state_dict(), PATH)

这里其实少了验证集评估，这样才能确定最佳训练的 epoch 轮数。

5 在测试数据上测试网络

我们已对训练数据集进行了两轮网络训练，但需要检查网络是否真的学到了任何内容。

我们将通过预测神经网络输出的类别标签，并将其与真实标签进行比对来检查。如果预测正确，就将该样本添加到正确预测列表中。

好的，第一步。让我们显示一个来自测试集的图像，以便熟悉一下。

dataiter = iter(testloader)
images, labels = next(dataiter)

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join(f'{classes[labels[j]]:5s}' for j in range(4)))

GroundTruth:  cat   ship  ship  plane

接下来，我们重新加载已保存的模型（注意：此处无需保存和重新加载模型，我们仅用于说明如何操作）：

net = Net()
net.load_state_dict(torch.load(PATH, weights_only=True))

<All keys matched successfully>

好的，现在让我们来看看神经网络对上述这些例子的判断：

outputs = net(images)

outputs：是网络 net(images) 的输出，形状为 [batch_size, 10]。对于刚才取的 4 张图片，它的形状是 [4, 10]，每一行存储的是模型认为该图片属于 10 个类别（0~9）的原始得分。

输出是10个类别的能量值。某个类别的能量越高，网络认为该图像属于该类别的可能性就越大。因此，我们来获取最高能量的索引：

_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join(f'{classes[predicted[j]]:5s}'
                              for j in range(4)))

Predicted:  cat   ship  car   plane

结果看起来相当不错。

让我们来看看该网络在整个数据集上的表现。

correct = 0
total = 0
net.eval()  # 切换到评估模式

# since we're not training, we don't need to calculate the gradients for our outputs
with torch.no_grad():
    for data in testloader:
        images, labels = data
        # calculate outputs by running images through the network
        outputs = net(images)
        # the class with the highest energy is what we choose as prediction
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

Accuracy of the network on the 10000 test images: 55 %

看起来比随机猜测要好得多，而随机猜测的准确率只有10%（从10个类别中随机选择一个）。似乎网络学到了一些东西。

我们看看哪些类别表现良好，哪些类别表现不佳：

# prepare to count predictions for each class
correct_pred = {classname: 0 for classname in classes}
total_pred = {classname: 0 for classname in classes}
net.eval()  # 切换到评估模式

# again no gradients needed
with torch.no_grad():
    for data in testloader:
        images, labels = data
        outputs = net(images)
        _, predictions = torch.max(outputs, 1)
        # collect the correct predictions for each class
        for label, prediction in zip(labels, predictions):
            if label == prediction:
                correct_pred[classes[label]] += 1
            total_pred[classes[label]] += 1


# print accuracy for each class
for classname, correct_count in correct_pred.items():
    accuracy = 100 * float(correct_count) / total_pred[classname]
    print(f'Accuracy for class: {classname:5s} is {accuracy:.1f} %')

Accuracy for class: plane is 59.9 %
Accuracy for class: car   is 70.7 %
Accuracy for class: bird  is 28.8 %
Accuracy for class: cat   is 32.2 %
Accuracy for class: deer  is 53.8 %
Accuracy for class: dog   is 51.8 %
Accuracy for class: frog  is 63.3 %
Accuracy for class: horse is 67.3 %
Accuracy for class: ship  is 65.5 %
Accuracy for class: truck is 66.4 %

之后我没看如何在 GPU 上运行这些神经网络模型

在GPU上训练

就像将张量转移到 GPU 上一样，你也需要将神经网络转移到 GPU 上。但是你需要确保模型和数据均在同一个设备上，不然就会报错。

首先，如果可用 CUDA，则将我们的设备定义为第一个可见的 CUDA 设备：

device = torch.device(torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu')
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 旧的查找 cuda 代码

# Assuming that we are on a CUDA machine, this should print a CUDA device:

print(device)

cuda

本节其余部分假设该设备为CUDA设备。

然后这些方法将递归遍历所有模块，并将它们的参数和缓冲区转换为CUDA张量：

net.to(device)

其实正常一开始构建好模型后，就可以直接转移到 GPU 上，代码如下。然后就可以在 GPU 上训练模型了。

net = Net().to(device)

并且数据也要发送到GPU：

inputs, labels = data[0].to(device), data[1].to(device)

为什么我没有注意到与 CPU 相比的大规模加速？因为你的网络真的很小。

AI 给的完整代码如下

import torch

# 查找设备
device = torch.device(torch.accelerator.current_accelerator().type if torch.accelerator.is_available() else 'cpu')
print(device)

# 加载网络
net = Net()
net.load_state_dict(torch.load(PATH, weights_only=True)) # 加载训练好的权重
net.to(device) # 放到 GPU 中

net.eval()  # 1. 切换到评估模式[reference:0][reference:1]
correct = 0
total = 0

with torch.no_grad():  # 2. 禁用梯度计算，节省内存和计算[reference:2]
    for images, labels in testloader:
        images, labels = images.to(device), labels.to(device) # 放到 GPU 中
        outputs = net(images)                # 3. 前向传播，得到预测logits
        _, predicted = torch.max(outputs, 1)   # 4. 取出概率最高的类别索引[reference:4]
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f'测试集准确率: {accuracy:.2f}%')        # 5. 计算并打印准确率[reference:6]

代码解释：

• net.eval()：将模型切换到评估模式，这对于包含 Dropout 或 BatchNorm 的模型至关重要。
• with torch.no_grad():：在该上下文管理器中，PyTorch不会计算梯度，能显著减少内存消耗并加速计算。

如果之后要让模型重新训练，记得切换回训练模式

net.train() 

参考资料

1. https://docs.pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html
2. https://www.bilibili.com/video/BV1ir1WBrEuz?vd_source=584d44ba7fed271fb82848e7894a53c8&spm_id_from=333.788.player.switch&p=2