torch.nn到底是什么？

关注

翻译最新推荐文章于 2026-06-05 14:02:43 发布 2.1w 阅读 AI 写同款· GEO 优化 ›

PyTorch 中神经网络库torch.nn的详细介绍 torch.nn是 PyTorch 深度学习框架中的一个核心模块，它为构建和训练神经网络提供了丰富的类库。通过torch.nn，开发者能够快速构建复杂的深度学习模型，并利用 PyTorch 动态计算图特性进行高效训练和推理。此外，该模块还与torch.optim配合，方便地进行权重优化；以及与DataLoader结合以组织和迭代训练数据。阅读详情

精简版可参考《torch.nn到底是什么？（精简版）》

PyTorch 提供了设计优雅的模块和类：torch.nn，torch.optim，Dateset 和 DataLoader，以帮助你创建和训练神经网络。为了充分利用它们的功能并且为你的问题定制它们，你需要正真理解它们在做什么。为了逐渐理解，我们首先在 MNIST 数据集上训练基本的神经网络，而不使用这些模块的任何特征。最初只会使用最基本的 PyTorch tensor 功能。然后，我们逐步添加来自 torch.nn，torch.optim，Dataset 和 DataLoader 的一个特征，以显示每一部分的功能，以及它如何使得代码更简洁或灵活。

1、设置MNIST数据

我们将使用经典的 MNIST 数据集，该数据集由手写数字（0-9）的黑白图像组成。

我们将使用 pathlib 来处理路径（Python3标准库的一部分），用 requests 下载数据。只有当我们需要模块的时候才会导入它们，因此你可以清楚地看到正在使用的模块。

from pathlib import Path
import requests

DATA_PATH = Path("data")
PATH = DATA_PATH / "mnist"

PATH.mkdir(parents=True, exist_ok=True)

URL = "http://deeplearning.net/data/mnist/"
FILENAME = "mnist.pkl.gz"

if not (PATH / FILENAME).exists():
        content = requests.get(URL + FILENAME).content
        (PATH / FILENAME).open("wb").write(content)

该数据集的格式为NumPy array，使用 pickle 存储。

import pickle
import gzip

with gzip.open((PATH / FILENAME).as_posix(), "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

每个图片大小为28x28，并存储为长度为784（=28x28）的扁平行。

查看其中的一个图片：

from matplotlib import pyplot
import numpy as np

pyplot.imshow(x_train[0].reshape((28, 28)), cmap="gray")
print(x_train.shape)

输出为：
在这里插入图片描述

(50000, 784)

PyTorch使用 tensor 而不是 NumPy array，所以我们需要将其转换。

import torch

x_train, y_train, x_valid, y_valid = map(
    torch.tensor, (x_train, y_train, x_valid, y_valid)
)
n, c = x_train.shape
x_train, x_train.shape, y_train.min(), y_train.max()
print(x_train, y_train)
print(x_train.shape)
print(y_train.min(), y_train.max())

输出：

tensor([[0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        ...,
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.],
        [0., 0., 0.,  ..., 0., 0., 0.]]) tensor([5, 0, 4,  ..., 8, 4, 8])
torch.Size([50000, 784])
tensor(0) tensor(9)

2、从头构建神经网络（不使用 `torch.nn`）

我们首先只使用PyTorch tensor 操作创建一个模型。PyTorch提供了创建随机tensor或零填充tensor的方法，我们将使用这些方法为简单线性模型创建权重（weight）和偏置值（bias）。它们都是普通的tensor，除此之外，我们增加了一点：我们告诉PyTorch它们需要梯度。这将使得PyTorch记录tensor上的所有操作，因此PyTorch可以在反向传播中自动计算梯度。

对于权重，我们在初始化之后设置 requires_grad，因为我们不想在梯度中包含这一步。

#initializing the weights with Xavier initialisation (by multiplying with 1/sqrt(n)).

import math

weights = torch.randn(784, 10) / math.sqrt(784)
weights.requires_grad_()
bias = torch.zeros(10, requires_grad=True)

由于PyTorch可以自动计算梯度，我们可以使用任何标准Python函数（或可调用对象）作为模型。因此我们只编写一个矩阵乘法和广播加法来创建一个简单的线性模型。我们还需要一个激活函数，因此我们将编写 log_softmax 并使用它。记住：虽然PyTorch提供了许多的预先编写好的损失函数、激活函数等，但是你可以使用普通的Python编写自己的函数。PyTorch甚至可以自动为你的函数创建快速GPU或矢量化CPU代码。

def log_softmax(x):
    return x - x.exp().sum(-1).log().unsqueeze(-1)

def model(xb):
    return log_softmax(xb @ weights + bias)

上面的代码中，"@"代表点积操作。我们将会在一批数据（64个图片）上调用我们编写的函数。这是一个前向传播。注意，现阶段我们的预测情况比随机猜测好不到哪里，因为我们是从随机权重开始的。

bs = 64                  # batch size

xb = x_train[0:bs]       # a mini-batch from x
preds = model(xb)        # predictions
preds[0], preds.shape
print(preds[0], preds.shape)

输出：

tensor([-1.7022, -3.0342, -2.4138, -2.6452, -2.7764, -2.0892, -2.2945, -2.5480,
        -2.3732, -1.8915], grad_fn=<SelectBackward>) torch.Size([64, 10])

如你所见，pred tensor不只包含tensor值，还包含一个梯度函数，我们随后将会使用它来来进行反向传播。

我们编写负对数似然函数并将其用作损失函数（我们可以只使用标准的Python实现）：

def nll(input, target):
    return -input[range(target.shape[0]), target].mean()

loss_func = nll

查看随机初始化模型的损失，这样我们随后就可以看到在使用了反向传播后，是否改进了模型。

yb = y_train[0:bs]
print(loss_func(preds, yb))

输出：

tensor(2.3783, grad_fn=<NegBackward>)

编写函数计算模型的精确度。对于每一次预测，如果最大值的索引和目标值匹配，则表示预测正确。

def accuracy(out, yb):
    preds = torch.argmax(out, dim=1)
    return (preds == yb).float().mean()

查看随机初始化模型的精确度，因此我们可以观察随着损失的改善，精确度是否提升。

print(accuracy(preds, yb))

输出：

tensor(0.0938)

现在我们可以进行训练。对于每次迭代，将会做以下几件事：

选择一批数据（mini-batch）
使用模型进行预测
计算损失
loss.backward() 更新模型的梯度，即权重和偏置

我们现在使用这些梯度来更新权重和偏置。我们将在 torch.no_grad() 中执行，因为
我们不想记录这些操作来进行下一次梯度计算。

我们随后将梯度设置为0，以便为下一次循环做好准备。否则，梯度将会记录所有发生的操作。（也就是说，loss.backward() 将梯度增加到已经存在的值上，而不是替代它）

from IPython.core.debugger import set_trace

lr = 0.5  # learning rate
epochs = 2  # how many epochs to train for

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
       #set_trace()
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()

到此，我们已经从头编写并训练了一个小型的神经网络。

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

输出：

tensor(0.0806, grad_fn=<NegBackward>) tensor(1.)

将损失和精确度与前边的比较，发现损失减少，精确度提升。

3、使用 `torch.nn.functional`

我们现在来重构代码，代码的功能和前边的一样，我们只是利用PyTorch 的 nn 类来使得代码更简洁和灵活。

第一步并且最简单的一步是用 torch.nn.functional（通常导入到命名空间F中）中的函数替代我们手工编写的激活函数和损失函数来缩短代码。该模块包含 torch.nn 库中所有的函数（而库的其他部分还包含类）。除了各种损失函数和激活函数，在还模块中你还可以发现许多用于创建神经网络的方便的函数，如池化函数等。

如果使用了负对数似然损失函数和 log softnax 激活函数，那么Pytorch提供的F.cross_entropy 结合了两者。所以我们甚至可以从我们的模型中移除激活函数。

import torch.nn.functional as F

loss_func = F.cross_entropy

def model(xb):
    return xb @ weights + bias

注意，在 model 函数中我们不再需要调用 log_softmax。让我们确认一下，损失和精确度与前边计算的一样：

print(loss_func(model(xb), yb), accuracy(model(xb), yb))

输出：

tensor(0.0806, grad_fn=<NllLossBackward>) tensor(1.)

4、使用 `nn.Module` 重构

下一步，我们将使用 nn.Module 和 nn.Parameter,以获得更清晰更简洁的训练循环。我们继承 nn.Module（它本身是一个类并且能够跟踪状态）建立子类。我们想要建立一个包含权重、偏置和前向传播的方法的类。nn.Module 拥有许多我们将会使用的属性和方法（例如：.parameters() 和.zero_grad()）。

from torch import nn

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(784, 10) / math.sqrt(784))
        self.bias = nn.Parameter(torch.zeros(10))

    def forward(self, xb):
        return xb @ self.weights + self.bias

因为我们现在使用一个对象而不是一个函数，所以我们首先需要实例化我们的模型：

model = Mnist_Logistic()

现在我们可以像之前那样计算损失。注意，nn.Module 对象像是函数一样被使用（即它们能被调用），但在幕后，PyTorch 将自动调用我们的 forward 方法。

print(loss_func(model(xb), yb))

输出：

tensor(2.3558, grad_fn=<NllLossBackward>)

以前对于我们的训练循环，我们需要按名字更新每个参数的值，并且手动将每个参数的梯度归零，像下面这样：

with torch.no_grad():
    weights -= weights.grad * lr
    bias -= bias.grad * lr
    weights.grad.zero_()
    bias.grad.zero_()

现在我们可以利用 model.paremeters() 和 model.zero_grad() 使得这些步骤更简洁，特别是，当我们有一个更复杂的模型时，使得我们更不容忘记某些参数。

with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()

我们将训练循环包装到一个 fit 函数中，以便我们以后运行。

def fit():
    for epoch in range(epochs):
        for i in range((n - 1) // bs + 1):
            start_i = i * bs
            end_i = start_i + bs
            xb = x_train[start_i:end_i]
            yb = y_train[start_i:end_i]
            pred = model(xb)
            loss = loss_func(pred, yb)

            loss.backward()
            with torch.no_grad():
                for p in model.parameters():
                    p -= p.grad * lr
                model.zero_grad()

fit()

再次检查损失是否下降

print(loss_func(model(xb), yb))

输出：

tensor(0.0826, grad_fn=<NllLossBackward>)

5、使用 `nn.Linear` 重构

继续重构代码。我们将会使用PyTorch 的 nn.Linear 类建立一个线性层，以替代手动定义和初始化 self.weights 和 self.bias、计算 xb @ self.weights + self.bias 等工作。PyTorch拥有多中类型预先定义好的层可以帮助我们极大简化代码，并且通常可以使之运行更快。

class Mnist_Logistic(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin = nn.Linear(784, 10)

    def forward(self, xb):
        return self.lin(xb)

我们像之前一样实例化模型并且计算损失

model = Mnist_Logistic()
print(loss_func(model(xb), yb))

输出：

tensor(2.3156, grad_fn=<NllLossBackward>)

我们仍然能够像之前那样使用 fit 方法

fit()

print(loss_func(model(xb), yb))

输出：

tensor(0.0809, grad_fn=<NllLossBackward>)

6、使用 `optim` 重构

PyTorch还有一个包含各种优化算法的包 torch.optim 。我们可以使用优化器中的 step 方法来执行前向步骤，而不是手动更新参数。

这将使得我们替换之前手动编写的优化步骤：

with torch.no_grad():
    for p in model.parameters(): p -= p.grad * lr
    model.zero_grad()

替换为：

opt.step()
opt.zero_grad()

optim.zero_grad() 将梯度重置为0，我们需要在计算下一个minibatch的梯度前调用它。

from torch import optim

我们将要定义一个函数来创建模型和优化器，以便将来可以重用它。

def get_model():
    model = Mnist_Logistic()
    return model, optim.SGD(model.parameters(), lr=lr)

model, opt = get_model()
print(loss_func(model(xb), yb))

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

输出：

tensor(2.2861, grad_fn=<NllLossBackward>)
tensor(0.0815, grad_fn=<NllLossBackward>)

7、使用 `Dataset` 重构

PyTorch 有一个抽象的 Dataset 类。任何具有 __len__（通过Python的标准len函数调用）函数和 __getitem__ 函数的类都可以是一个 Dataset。

Pytorch 的 TensorDataset 是一个包装 tensor 的 Dataset。通过定义长度和索引方式，这也为我们提供了一种沿tensor第一维迭代、索引、切片的方法。这将使我们更容易在我们训练的同一行中访问独立变量和因变量。

from torch.utils.data import TensorDataset

x_train 和y_train 可以组合在一个单独的 TensorDataset 中，这将更容易迭代和切片。

train_ds = TensorDataset(x_train, y_train)

之前，我们必须分别迭代x和y的小批量值：

xb = x_train[start_i:end_i]
yb = y_train[start_i:end_i]

现在，我们可以一起做这两步：

xb,yb = train_ds[i*bs : i*bs+bs]

model, opt = get_model()

for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        xb, yb = train_ds[i * bs: i * bs + bs]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

输出：

tensor(0.0800, grad_fn=<NllLossBackward>)

8、使用 `DataLoader` 重构

PyTorch 的 DataLoader 负责管理批次。你可以从任何 Dataset 创建 DataLoader。DataLoader 使得迭代批次更简单。DataLoader 自动地提供每个批次，而不必使用train_ds[i*bs : i*bs+bs]。

from torch.utils.data import DataLoader

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs)

之前，我们像下面这样迭代批次：

for i in range((n-1)//bs + 1):
    xb,yb = train_ds[i*bs : i*bs+bs]
    pred = model(xb)

现在，我们的循环更加简洁，因为（xb，yb）可以自动从 DataLoader 自动加载：

for xb,yb in train_dl:
    pred = model(xb)
model, opt = get_model()

for epoch in range(epochs):
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

print(loss_func(model(xb), yb))

输出：

tensor(0.0821, grad_fn=<NllLossBackward>)

由于PyTorch的 nn.Module、nn.Parameter、Dataset 和 DataLoader，现在我们的训练循环变得更小、更容易理解。现在让我们尝试添加在实践中创建有效模型所需的基本功能。

9、增加验证

在第1部分中，我们只是尝试去设置合理的训练循环以用于我们的训练数据。实际上，你总是应该有一个验证集，以确定你是否过度拟合。

打乱训练数据对于防止批次与过度拟合之间的相关性非常重要。另一方面，无论我们是否打乱验证集，验证损失都是相同的。由于打乱需要额外的时间，因此打乱验证数据是没有意义的。

我们将要使用的验证集的大小是训练集的两倍。这是因为验证集不需要反向传播，因此占用更少的内存（不需要存储梯度）。我们利用这一点来使用更大的批次大小并更快地计算损失。

train_ds = TensorDataset(x_train, y_train)
train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True)

valid_ds = TensorDataset(x_valid, y_valid)
valid_dl = DataLoader(valid_ds, batch_size=bs * 2)

我们将在每个epoch结束时计算和打印验证损失。（注意，我们总是在训练之前调用model.train()，在推理之前调用 model.eval()，因为这些由诸如 nn.BatchNorm2d 和nn.Dropout 等层使用，以确保这些不同阶段的适当行为。）

model, opt = get_model()

for epoch in range(epochs):
    model.train()
    for xb, yb in train_dl:
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        opt.step()
        opt.zero_grad()

    model.eval()
    with torch.no_grad():
        valid_loss = sum(loss_func(model(xb), yb) for xb, yb in valid_dl)

    print(epoch, valid_loss / len(valid_dl))

输出：

0 tensor(0.2981)
1 tensor(0.3033)

10、创建 `fit()` 和 `get_data()`

我们现在将进行一些重构。因为计算训练集和验证集的损失，我们进行了两次相似的处理，让我们将其作为一个 loss_batch 函数来计算每个批次的损失。

我们为训练集传递一个优化器，并使用它来执行反向传播。对于验证集，我们不传递优化器，因此该方法不执行反向传播。

def loss_batch(model, loss_func, xb, yb, opt=None):
    loss = loss_func(model(xb), yb)

    if opt is not None:
        loss.backward()
        opt.step()
        opt.zero_grad()

    return loss.item(), len(xb)

fit 运行必要的操作来训练我们的模型并计算每个epoch的训练和验证损失。

import numpy as np

def fit(epochs, model, loss_func, opt, train_dl, valid_dl):
    for epoch in range(epochs):
        model.train()
        for xb, yb in train_dl:
            loss_batch(model, loss_func, xb, yb, opt)

        model.eval()
        with torch.no_grad():
            losses, nums = zip(
                *[loss_batch(model, loss_func, xb, yb) for xb, yb in valid_dl]
            )
        val_loss = np.sum(np.multiply(losses, nums)) / np.sum(nums)

        print(epoch, val_loss)

get_data 为训练集合验证集返回 DataLoader。

def get_data(train_ds, valid_ds, bs):
    return (
        DataLoader(train_ds, batch_size=bs, shuffle=True),
        DataLoader(valid_ds, batch_size=bs * 2),
    )

现在，我们获取 DataLoader 和拟合模型的整个过程可以在3行代码中运行：

train_dl, valid_dl = get_data(train_ds, valid_ds, bs)
model, opt = get_model()
fit(epochs, model, loss_func, opt, train_dl, valid_dl)

输出：

0 0.3055081913471222
1 0.31777948439121245

11、总结

我们现在有一个通用数据流水线和训练循环，你可以使用它来训练多种类型PyTorch模型。各部分的功能总结如下：

torch.nn
- Module：创建一个可调用的对象，其行为类似于一个函数，但也可以包含状态（例如神经网络层权重）。它知道它包含哪些参数，并且可以将所有梯度归零，循环遍历它们更新权重等。
- Parameter：tensor 的包装器（wrapper），它告诉 Module 它具有在反向传播期间需要更新的权重。只更新具有 requires_grad 属性的 tensor。
- functional：一个模块（通常按惯例导入到F命名空间中），它包含激活函数，损失函数等，以及非状态（non-stateful）版本的层，如卷积层和线性层。
torch.optim：包含 SGD 等优化器，可在后向传播步骤中更新 Parameter 的权重。
Dataset：带有 __len__ 和 __getitem__ 的对象的抽象接口，包括 PyTorch 提供的类，如TensorDataset。
DataLoader：获取任何 Dataset 并创建一个返回批量数据的迭代器。