《动手学深度学习》（d2l-zh）：深度学习计算核心组件实战——层与块、参数管理、自定义层、模型读写与 GPU 加速

• 人工智能
• 深度学习
• 机器学习
• 教程

导读

本章是《动手学深度学习》中文版（d2l-zh）中承上启下的关键章节：前面章节已经用高级 API 组装过完整模型，但跳过了框架底层的工作细节；本章则深入探究深度学习计算的核心组件——模型构建（层与块）、参数访问与初始化、延后初始化、自定义层与块、模型读写到磁盘，以及利用 GPU 实现显著加速。读完本章，读者将从深度学习"基础用户"进阶为"高级用户"，在后续卷积网络、循环网络等高级模型章节中得以自如地设计、调试与部署自己的架构。

本仓库的章节正文位于 chapter_deep-learning-computation 目录，其中 index.md 是本章导言，随后由model-construction（层和块）、parameters（参数管理）、deferred-init（延后初始化）、custom-layer（自定义层）、read-write（读写文件）、use-gpu（GPU）六个小节组成，本篇文章即沿着这条主线展开。

为什么深度学习库的抽象在变"粗糙"？

从 2007 年发布的开创性 Theano 库开始，灵活的开源工具使研究人员能够快速开发模型原型，避免重复造轮子，同时仍保留底层修改的能力。随着时间的推移，深度学习库逐渐演化为提供越来越粗糙（coarse）的抽象：

• 半导体设计师从指定晶体管，到设计逻辑电路，再到编写代码；
• 神经网络研究人员同样从考虑单个神经元的行为，转变为从层的角度构思网络，如今在设计架构时往往直接以更粗糙的"块"（block）为单位。

以 ResNet-152 为例，这个赢得 2015 年 ImageNet 和 COCO 竞赛识别与检测任务的架构拥有数百层，而这些层其实是由若干"层组"（groups of layers）的重复模式组成的；在自然语言处理、语音等领域，类似的重复模式架构同样普遍。这种模块化的设计思想，正是本章引入"块"这一抽象的根本动机。

层和块：一切模型的通用抽象

从神经元、层到块

回顾多层感知机的演进：单个神经元接收输入、生成标量输出并带有一组可更新参数；层是一组神经元的矢量化描述；而整个多层感知机模型，连同其每一层，都符合"接收输入、生成输出、带有一组参数"的同一模式。研究讨论"比单个层大、比整个模型小"的组件（即层组）更有价值——块（block）正是这种抽象。

块的递归组合：一个块可以描述单个层、由多个层组成的组件，甚至整个模型本身；块又可以组合成更大的组件，这一过程通常是递归的。通过定义代码按需生成任意复杂度的块，我们可以用简洁代码实现复杂网络。

从编程角度看，块由类（class）表示，任何子类都必须：

1. 定义一个将输入转换为输出的前向传播函数；
2. 存储任何必需的参数（有些块不需要参数）；
3. 具备反向传播能力（由于自动微分提供了后端实现，自定义块时只需关注前向传播和参数）。

Sequential：最简单的块组合

回顾多层感知机的简洁实现：一个含 256 个单元、ReLU 激活的隐藏层，加上一个 10 单元无激活的输出层。以 PyTorch 为例：

import torch from torch import nn from torch.nn import functional as F net = nn.Sequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10)) X = torch.rand(2, 20) net(X)

各框架中的Sequential本质相同：MXNet 的nn.Sequential是维护有序Block列表的特殊Block；PyTorch 的nn.Sequential是维护有序Module列表的特殊Module；TensorFlow 的Sequential是特殊的keras.Model；PaddlePaddle 的nn.Sequential是维护有序Layer列表的特殊Layer。前向传播时，它把列表中的每个块依次串联，将上一个块的输出作为下一个块的输入。net(X)不过是net.__call__(X)的简写。

自定义块：MLP 从头实现

块的基本职责可总结为五点：将输入作为前向传播函数的参数；通过前向传播生成输出（输出形状可以不同于输入，例如上面的隐藏层将 20 维输入映射为 256 维输出）；计算输出关于输入的梯度（通常自动发生）；存储和访问前向传播所需的参数；根据需要初始化模型参数。

下面从零编写一个含 256 个隐藏单元、10 维输出的 MLP 块：

class MLP(nn.Module): def __init__(self): super().__init__() self.hidden = nn.Linear(20, 256) # 隐藏层 self.out = nn.Linear(256, 10) # 输出层 def forward(self, X): return self.out(F.relu(self.hidden(X))) net = MLP() net(X)

关键细节：

• 自定义__init__通过super().__init__()调用父类构造函数，省去模板代码；
• 两个全连接层作为实例变量self.hidden与self.out保存，实例化多个 MLP 时各自拥有独立的层对象，因此可以学习到不同模型；
• 除非实现新算子，反向传播函数与参数初始化都由框架自动生成。

块的另一大优势是通用性：可以子类化出层（如全连接层类）、整个模型（如MLP类）或各种中等复杂度的组件——后续卷积网络章节将充分用到这种灵活性。

亲手实现 MySequential：理解顺序块内部机制

Sequential的本质只有两件事：把块逐个追加到列表；前向传播时按追加顺序把输入依次传给每个块。以 PyTorch 的实现为例：

class MySequential(nn.Module): def __init__(self, *args): super().__init__() for idx, module in enumerate(args): # _modules 是有序字典（OrderedDict） self._modules[str(idx)] = module def forward(self, X): for block in self._modules.values(): X = block(X) return X net = MySequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10)) net(X)

各框架都把子块存放在专属的有序字典中（MXNet 的_children、PyTorch 的_modules、Paddle 的_sub_layers）。这样做的关键原因是：在参数初始化过程中，框架知道要到这些字典中查找需要初始化参数的子块——这是使用有序字典而非普通 Python 列表的核心优势。这也正是本章练习中"如果将存储块的方式改为 Python 列表会出现什么问题"的答案所在。

在前向传播中执行任意代码：常数参数与控制流

并非所有架构都是简单的顺序架构。块允许我们在前向传播函数中执行 Python 控制流与任意数学运算，并引入"常数参数"（constant parameter）——既非上一层结果、也非可更新参数的项。

例如计算 $f(\mathbf{x},\mathbf{w}) = c \cdot \mathbf{w}^\top \mathbf{x}$，其中 $c$ 是优化过程中不更新的常量。FixedHiddenMLP展示了这种做法：隐藏层权重在实例化时随机初始化后即固定不变，不参与反向传播更新；随后输出经过一个全连接层；返回前还运行 while 循环，在 $L_1$ 范数大于 1 时反复将输出除以 2，最后返回所有项的和。

class FixedHiddenMLP(nn.Module): def __init__(self): super().__init__() # 不计算梯度的随机权重参数，训练期间保持不变 self.rand_weight = torch.rand((20, 20), requires_grad=False) self.linear = nn.Linear(20, 20) def forward(self, X): X = self.linear(X) X = F.relu(torch.mm(X, self.rand_weight) + 1) X = self.linear(X) # 复用全连接层：相当于两全连接层共享参数 while X.abs().sum() > 1: # 控制流 X /= 2 return X.sum()

这些操作未必会在实际任务中使用，其意义在于展示"如何将任意代码集成到神经网络计算的流程中"。更进一步，块可以自由嵌套混合——例如NestMLP内部先用nn.Sequential组合若干层，再叠加自己的全连接层；还可以构造chimera = nn.Sequential(NestMLP(), nn.Linear(16, 20), FixedHiddenMLP())这样的"嵌合体"网络。

效率的考量

在一个高性能深度学习库中大量执行 Python 字典查找与控制流，会触发 Python 全局解释器锁（GIL）的开销问题：速度极快的 GPU 可能要等待 CPU 跑完 Python 代码才能继续作业。MXNet 的 Gluon 通过"混合式编程"（hybridization）在首次调用时记录计算图、后续对调用进行加速（详见 chapter_computational-performance/hybridize.md），但当控制流在不同网络通路引导不同分支时需格外小心。

参数管理：访问、初始化与绑定

训练的目标是找到使损失函数最小化的参数值；训练后要用参数做预测、复用甚至部署，因此必须掌握参数的访问、初始化和共享。

访问参数

以一个单隐藏层 MLP 为例（隐藏层 4→8，输出层 8→1）：

import torch from torch import nn net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1)) X = torch.rand(size=(2, 4)) net(X)

逐层访问：net[2].state_dict()显示输出层包含权重与偏置两个参数，均以单精度浮点数（float32）存储，参数名允许在网络中唯一标识每个参数。访问具体数值需先取出参数对象再取底层数据：

print(type(net[2].bias)) print(net[2].bias) print(net[2].bias.data)

参数是复合对象，包含值、梯度和额外信息——这正是需要显式取.data的原因。未调用反向传播时梯度处于初始状态，可用net[2].weight.grad == None验证。

一次性访问所有参数：嵌套块结构下逐个访问十分麻烦，需要递归整棵树。可以用named_parameters()递归列出全部参数：

print(*[(name, param.shape) for name, param in net.named_parameters()])

或者通过net.state_dict()['2.bias'].data按参数字典键名访问。

从嵌套块收集参数：用"块工厂"函数构造嵌套网络——block1()生成两层 Sequential，block2()循环 4 次嵌套block1()，再组合成rgnet = nn.Sequential(block2(), nn.Linear(4, 1))。打印rgnet可以看到清晰的层级结构；由于层是分层嵌套的，可以像嵌套列表一样索引，例如rgnet[0][1][0].bias.data访问第一个主块中第二个子块第一层的偏置。

参数初始化

各框架默认初始化策略不同：MXNet 默认从 $U(-0.07, 0.07)$ 采样权重并将偏置置 0；PyTorch 按输入输出维度计算的范围均匀初始化；Keras 按输入输出维度范围均匀初始化权重、偏置置 0；PaddlePaddle 默认使用 Xavier 初始化。框架也都提供自定义初始化入口。

内置初始化：将所有权重初始化为标准差 0.01 的高斯随机变量、偏置为 0：

def init_normal(m): if type(m) == nn.Linear: nn.init.normal_(m.weight, mean=0, std=0.01) nn.init.zeros_(m.bias) net.apply(init_normal) net[0].weight.data[0], net[0].bias.data[0]

初始化为常数 1 只需把nn.init.normal_换成nn.init.constant_(m.weight, 1)。不同层使用不同初始化也很简单：对第一个线性层应用nn.init.xavier_uniform_，对输出层应用nn.init.constant_(m.weight, 42)，用net[0].apply(init_xavier)与net[2].apply(init_42)分别作用即可。

自定义初始化：框架没有所需方法时，可以自行定义。例如希望权重 $w$ 以 1/4 概率取自 $U(5,10)$、以 1/2 概率为 0、以 1/4 概率取自 $U(-10,-5)$：

def my_init(m): if type(m) == nn.Linear: print("Init", *[(name, param.shape) for name, param in m.named_parameters()][0]) nn.init.uniform_(m.weight, -10, 10) m.weight.data *= m.weight.data.abs() >= 5 net.apply(my_init) net[0].weight[:2]

此外也可以直接就地修改参数值，例如net[0].weight.data[:] += 1、net[0].weight.data[0, 0] = 42。MXNet 高级用户需注意：若要在autograd范围内调整参数，应使用set_data以避免误导自动微分机制。

参数绑定（共享参数）

有时需要在多个层间共享参数。PyTorch 中直接把同一个Linear实例多次放入网络即可：

shared = nn.Linear(8, 8) net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), shared, nn.ReLU(), shared, nn.ReLU(), nn.Linear(8, 1)) net(X) print(net[2].weight.data[0] == net[4].weight.data[0]) # True net[2].weight.data[0, 0] = 100 print(net[2].weight.data[0] == net[4].weight.data[0]) # True，同一张量

参数绑定意味着两个位置不仅值相等，而且由同一个张量表示，修改其一、另一随之改变。绑定带来的梯度行为是：由于模型参数包含梯度，反向传播时两个绑定位置的梯度会累加。MXNet 中则通过nn.Dense(8, activation='relu', params=shared.params)显式传入共享参数实现；TensorFlow 的 Keras 略有不同——它会自动去重重复层（上述例子打印len(net.layers) == 3）。

延后初始化：不知道输入维度也能建模型

前面构建网络时其实"违反"了三件事：定义网络时未指定输入维度；添加层时未指定前一层的输出维度；初始化参数时甚至不知道模型该有多少参数。代码能运行的关键是延后初始化（deferred initialization）：框架直到数据第一次通过模型时，才动态推断每个层的大小。

对卷积神经网络尤其方便——输入分辨率会层层影响后续维数，无需事先知道维度即可设置参数，大大简化模型定义与修改。

net = tf.keras.models.Sequential([ tf.keras.layers.Dense(256, activation=tf.nn.relu), tf.keras.layers.Dense(10), ]) # 此时各层对象存在，但权重为空 [net.layers[i].get_weights() for i in range(len(net.layers))]

MXNet 中此时每个层的输入维度显示为特殊值 -1（表示维度未知），直接访问net[0].weight.data()会触发"必须先初始化"的运行时错误；调用net.initialize()后一切照旧——它只是声明"希望初始化"并可选指定分布，并不会真正初始化。

直到数据通过模型：

X = tf.random.uniform((2, 20)) net(X) [w.shape for w in net.get_weights()]

一旦知道输入维数为 20，框架即可确定第一层权重矩阵形状，然后依次处理第二层、第三层……直到所有形状已知后再真正初始化参数。通常只有第一层需要延后初始化，但框架仍是按顺序依次初始化。

自定义层：从无参数层到带参数层

神经网络的灵活性是深度学习成功的关键因素之一；当现有框架缺少所需层时，就必须构建自定义层。

不带参数的层

继承基础层类并实现前向传播即可。下面的CenteredLayer从输入中减去均值：

import torch from torch import nn import torch.nn.functional as F class CenteredLayer(nn.Module): def __init__(self): super().__init__() def forward(self, X): return X - X.mean() layer = CenteredLayer() layer(torch.FloatTensor([1, 2, 3, 4, 5]))

验证结果[-2., -1., 0., 1., 2.]符合预期。该层还能作为组件嵌入更复杂模型：net = nn.Sequential(nn.Linear(8, 128), CenteredLayer())，向网络发送随机数据后检查Y.mean()，由于浮点存储精度，会看到极小的非零值（接近 0）。

带参数的层

可训练参数需要用内置函数创建，以自动获得访问、初始化、共享、保存与加载等管理功能——好处是不必为每个自定义层编写自定义序列化程序。自定义全连接层需要权重与偏置两个参数，输入参数in_units（输入数）与units（输出数），使用 ReLU 激活：

class MyLinear(nn.Module): def __init__(self, in_units, units): super().__init__() self.weight = nn.Parameter(torch.randn(in_units, units)) self.bias = nn.Parameter(torch.randn(units,)) def forward(self, X): linear = torch.matmul(X, self.weight.data) + self.bias.data return F.relu(linear) linear = MyLinear(5, 3) linear.weight

• MXNet 用self.params.get('weight', shape=(in_units, units))声明参数；
• TensorFlow 通过build(X_shape)方法配合add_weight在知道输入形状后创建参数；
• Paddle 用paddle.create_parameter(shape=..., dtype='float32')创建。

自定义层可以独立执行前向传播（如linear(torch.rand(2, 5))），也可以像内置层一样直接用于构建模型：

net = nn.Sequential(MyLinear(64, 8), MyLinear(8, 1)) net(torch.rand(2, 64))

读写文件：保存张量与模型参数

训练耗时较长时，最佳实践是定期保存中间结果，避免断电等意外损失数天计算成果；部署场景也需要把训练好的模型保存到磁盘。

保存与加载张量

单个张量可直接调用save/load：

import torch from torch import nn x = torch.arange(4) torch.save(x, 'x-file') x2 = torch.load('x-file')

张量列表与字典同样支持：torch.save([x, y], 'x-files')后x2, y2 = torch.load('x-files')；mydict = {'x': x, 'y': y}; torch.save(mydict, 'mydict')后读回——"字符串映射到张量"的字典读写方式，正好适用于后续保存模型的所有权重。TensorFlow 侧对应使用np.save('x-file.npy', x)/np.load(..., allow_pickle=True)。

保存与加载模型参数

重要细节：保存的是模型的参数（state_dict），而不是整个模型。因为模型本身可以包含任意代码，难以序列化；恢复模型时需要先用代码重建架构，再从磁盘加载参数。

class MLP(nn.Module): def __init__(self): super().__init__() self.hidden = nn.Linear(20, 256) self.output = nn.Linear(256, 10) def forward(self, x): return self.output(F.relu(self.hidden(x))) net = MLP() X = torch.randn(size=(2, 20)) Y = net(X) torch.save(net.state_dict(), 'mlp.params') # 保存参数 clone = MLP() # 重建架构 clone.load_state_dict(torch.load('mlp.params')) # 加载参数 clone.eval() Y_clone = clone(X) Y_clone == Y # True

各框架对应 API：MXNetnet.save_parameters('mlp.params')/load_parameters；TensorFlownet.save_weights('mlp.params')/load_weights；Paddlepaddle.save(net.state_dict(), 'mlp.pdparams')/set_state_dict(paddle.load(...))。由于两个实例参数相同，相同输入X下Y_clone == Y逐元素成立。

GPU 使用：把计算搬到显卡上

自 2000 年以来 GPU 性能约每十年增长 1000 倍，深度学习研究离不开 GPU。本节讨论单 GPU 用法，并延伸至多 GPU、多服务器场景。

准备工作与计算设备

首先确保至少安装了一块 NVIDIA GPU，安装 NVIDIA 驱动与 CUDA 并设置路径，然后运行nvidia-smi查看显卡信息。MXNet/Paddle 需确认安装的是对应 CUDA 版本的 GPU 版框架（如pip install mxnet-cu100对应 CUDA 10.0）。本节多 GPU 演示需要至少两个 GPU，云上的多 GPU 实例（如 AWS EC2 多 GPU 实例）很容易获得。

各框架中 CPU/GPU 的表示方式：

注意语义差异：cpu()表示所有物理 CPU 和内存（计算会尝试使用所有 CPU 核心），而gpu()只代表一块卡及相应显存；gpu(0)与gpu()等价。

查询可用 GPU 数量：PyTorchtorch.cuda.device_count()、MXNetnpx.num_gpus()、TensorFlowlen(tf.config.experimental.list_physical_devices('GPU'))、Paddlepaddle.device.cuda.device_count()。

在无 GPU 环境下也能运行的辅助函数try_gpu与try_all_gpus是本书高频工具，已在仓库的 d2l/torch.py、d2l/mxnet.py、d2l/tensorflow.py、d2l/paddle.py 中随书实现并保存，其核心逻辑即"存在 GPU 返回 GPU，否则回退 CPU"：

def try_gpu(i=0): """如果存在，则返回gpu(i)，否则返回cpu()""" if torch.cuda.device_count() >= i + 1: return torch.device(f'cuda:{i}') return torch.device('cpu') def try_all_gpus(): """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]""" devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())] return devices if devices else [torch.device('cpu')]

张量与 GPU

默认张量在 CPU 上创建（x.device/x.ctx/x.place可查询）。无论何时对多个项进行操作，它们必须在同一设备上，否则框架不知道结果存哪里、在哪里执行计算。在 GPU 上创建张量，创建时指定设备即可：

X = torch.ones(2, 3, device=try_gpu()) # 第一个 GPU 上的全 1 张量 Y = torch.rand(2, 3, device=try_gpu(1)) # 第二个 GPU 上的随机张量

GPU 上创建的张量只消耗该 GPU 显存，可用nvidia-smi查看显存占用，注意不要创建超出显存限制的数据。

设备间复制

计算X + Y前必须让二者位于同一设备。把X复制到第二个 GPU：

Z = X.cuda(1) print(X) # 仍位于第一个 GPU print(Z) # 已复制到第二个 GPU Y + Z # 现在同一设备，可以相加

复制行为细节：PyTorch 中若Z已位于目标设备，Z.cuda(1)会直接返回Z而不复制分配新内存（Z.cuda(1) is Z为 True）；MXNet 的copyto即使目标设备相同也会复制，as_in_ctx则只在设备不同时复制，通常更推荐。

性能旁注：设备间数据传输远慢于计算，频繁拷贝会拖慢整体并增加并行化难度。经验法则：多个小操作比一个大操作糟糕得多；一次执行几个操作优于代码中散布的许多单个操作；一个设备等待另一个设备会产生阻塞。打印张量或转 NumPy 时，若数据不在内存中，框架会先复制到内存——这还会引入全局解释器锁开销，使一切等待 Python 完成。因此应尽量在 GPU 内部为日志分配内存，只移动较大的日志。

神经网络与 GPU

模型同样可以指定设备。把模型参数放到 GPU 上：PyTorch 用net = net.to(device=try_gpu())，MXNet 用net.initialize(ctx=try_gpu())，TensorFlow 用tf.distribute.MirroredStrategy()的作用域创建模型，Paddle 用net = net.to(try_gpu())。

当输入是 GPU 上的张量时，模型会在同一 GPU 上计算结果；可进一步确认参数存储位置（net[0].weight.data.device）。只要所有数据和参数都在同一设备上，就能高效地学习模型——后续章节将有大量此类实例。

小结

• 一个块可以由许多层组成，也可以由许多块组成；块可以包含代码，并负责参数初始化与反向传播等大量内部处理。
• 层和块的顺序连接由Sequential块处理；自定义块只需实现构造函数与前向传播。
• 参数管理提供了访问、初始化与绑定等多种手段，支持自定义初始化方法。
• 延后初始化使框架能自动推断参数形状，简化模型架构修改。
• save/load函数可用于张量对象读写；参数字典可保存和加载网络全部参数，但架构必须用代码重建。
• 计算设备可显式指定（CPU/GPU），但所有输入数据必须在同一设备上；不经意地移动数据会显著降低性能。

进阶练习方向

1. 将MySequential的存储方式改为 Python 列表，思考参数初始化为何失效；
2. 实现一个"平行块"：接收两个块net1、net2，在前向传播中返回两个网络输出的串联；
3. 实现一个函数，生成同一块的多个实例并在此基础上构建更大的网络；
4. 基于 chapter_deep-learning-computation/parameters.md 中定义的FancyMLP访问各层参数，并观察共享参数层训练过程中的参数与梯度变化；
5. 实测大矩阵乘法的 CPU/GPU 速度差异，并验证"GPU 上保存日志、仅传输最终结果"与逐次传输的性能差距；
6. 尝试同时在两个 GPU 上执行两个矩阵乘法，与单 GPU 顺序执行对比，观察近线性的扩展效果。

本章各小节的完整代码与配套讨论均可在 chapter_deep-learning-computation 目录下的对应文档中找到，读者可按model-construction、parameters、deferred-init、custom-layer、read-write、use-gpu的顺序逐一研读与运行。

• 人工智能
• 深度学习
• 机器学习
• 教程