【动手学深度学习】第五章笔记：层与块、参数管理、自定义层、读写文件、GPU

为了更好的阅读体验，请点击这里

由于本章内容比较少且以后很显然会经常回来翻，因此会写得比较详细。

5.1 层和块

事实证明，研究讨论“比单个层大”但“比整个模型小”的组件更有价值。例如，在计算机视觉中广泛流行的ResNet-152 架构就有数百层，这些层是由层组（groups of layers）的重复模式组成。

为了实现这些复杂的网络，我们引入了神经网络块的概念。块（block）可以描述单个层、由多个层组成的组件或整个模型本身。使用块进行抽象的一个好处是可以将一些块组合成更大的组件。通过定义代码来按需生成任意复杂度的块，我们可以通过简洁的代码实现复杂的神经网络。

从编程的角度来看，块由类（class）表示。它的任何子类都必须定义一个将其输入转换为输出的前向传播函数，并且必须存储任何必需的参数。注意，有些块不需要任何参数。最后，为了计算梯度，块必须具有反向传播函数。在定义我们自己的块时，由于自动微分提供了一些后端实现，我们只需要考虑前向传播函数和必需的参数。

之后原书中举的例子为实例化一个包含两个线性层的多层感知机。该代码中，通过实例化 nn.Sequential 来构建模型，层的执行顺序是作为参数传递的。简而言之，nn.Sequential 定义了一种特殊的 Module，即在 PyTorch 中表示一个块的类，它维护了一个由 Module 组成的有序列表。注意，两个全连接层都是 Linear 类的实例，Linear 类本身就是 Module 的子类。另外，到目前为止，我们一直在通过 net(X) 调用我们的模型来获得模型的输出。这实际上是 net.__call__(X) 的简写。

5.1.1 自定义块

实现自定义块之前，简要总结一下每个块必须提供的基本功能。

1. 将输入数据作为其前向传播函数的参数。
2. 通过前向传播函数来生成输出。请注意，输出的形状可能与输入的形状不同。例如，我们上面模型中的第一个全连接的层接收一个20维的输入，但是返回一个维度为256的输出。
3. 计算其输出关于输入的梯度，可通过其反向传播函数进行访问。通常这是自动发生的。
4. 存储和访问前向传播计算所需的参数。
5. 根据需要初始化模型参数。

在下面的代码片段中，我们从零开始编写一个块。它包含一个多层感知机，其具有 \(256\) 个隐藏单元的隐藏层和一个 \(10\) 维输出层。注意，下面的 MLP 类继承了表示块的类。我们的实现只需要提供我们自己的构造函数（Python中的 __init__ 函数）和前向传播函数。

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

注意一些关键细节：首先，我们定制的 __init__ 函数通过 super().__init__() 调用父类的 __init__ 函数，省去了重复编写模版代码的痛苦。然后，我们实例化两个全连接层，分别为 self.hidden 和 self.out。注意，除非我们实现一个新的运算符，否则我们不必担心反向传播函数或参数初始化，系统将自动生成这些。

块的一个主要优点是它的多功能性。我们可以子类化块以创建层（如全连接层的类）、整个模型（如上面的MLP类）或具有中等复杂度的各种组件。

5.1.2 顺序块

构建简化的 MySequential，只需要定义两个关键函数：

1. 一种将块逐个追加到列表中的函数；
2. 一种前向传播函数，用于将输入按追加块的顺序传递给块组成的“链条”。

下面的 MySequential 类提供了与默认 Sequential 类相同的功能。

class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        for idx, module in enumerate(args):
            # 这里，module 是 Module 子类的一个实例。我们把它保存在 'Module' 类的成员
            # 变量 _modules 中。_module 的类型是 OrderedDict
            self._modules[str(idx)] = module

    def forward(self, X):
        # OrderedDict 保证了按照成员添加的顺序遍历它们
        for block in self._modules.values():
            X = block(X)
        return X

__init__ 函数将每个模块逐个添加到有序字典 _modules 中。读者可能会好奇为什么每个 Module 都有一个 _modules 属性？以及为什么我们使用它而不是自己定义一个Python列表？简而言之，_modules 的主要优点是：在模块的参数初始化过程中，系统知道在 _modules 字典中查找需要初始化参数的子块。

5.1.3 在前向传播函数中执行代码

当需要更强的灵活性时，我们需要定义自己的块。例如，可能希望在前向传播函数中执行Python的控制流。此外，可能希望执行任意的数学运算，而不是简单地依赖预定义的神经网络层。

那么，就可以在前向传播的函数中实现复杂的代码。

练习题

（1）如果将 MySequential 中存储块的方式更改为 Python 列表，会出现什么样的问题？

class MySequential(nn.Module):
    def __init__(self, *args):
        super().__init__()
        self.modules_list = []
        for idx, module in enumerate(args):
            self.modules_list.append(module)
        print(self.modules_list)
    
    def forward(self, X):
        for block in self.modules_list:
            X = block(X)
        return X
    
net = MySequential(nn.Linear(20, 256), nn.ReLU(), nn.Linear(256, 10))
net(X)

接下来如果调用 net.parameters() 迭代器来遍历参数或者用 net.state_dict() 来查看状态字典，你会发现什么也不会输出。原因在于 parameter 类型的参数只能从 _modules 中以及其他显示定义在表层的 nn.Module 类及子类获得，即使你把 list 换成另一个 OrderedDict 也并不好用。现在没办法自动获取了。

除此之外，由于无法自动获取 parameter 类型的参数，因此初始化很难做。

（2）实现一个块，它以两个块为参数，例如 net1 和 net2，并返回前向传播中两个网络的串联输出。这也被称为平行块。

class ParallelBlock(nn.Module):
    def __init__(self, net1, net2):
        super().__init__()
        self.net1 = net1
        self.net2 = net2
    
    def forward(self, X):
        return self.net2(self.net1(X))
    
net = ParallelBlock(nn.Linear(16, 20), nn.Linear(20, 10))
print(net)
for param in net.parameters():
    print(param)

（3）假设我们想要连接同一网络的多个实例。实现一个函数，该函数生成同一个块的多个实例，并在此基础上构建更大的网络。

一般而言 Sequential 就足够完成这个任务：

class multilayer(nn.Module):
    def __init__(self, num):
        super().__init__()
        layer_list = []
        for i in range(num):
            layer_list.append(nn.Linear(20, 10))
        self.ln = nn.Sequential(*layer_list)
    
    def forward(self, X):
        return self.ln(X)

multilayer(
  (ln): Sequential(
    (0): Linear(in_features=20, out_features=10, bias=True)
    (1): Linear(in_features=20, out_features=10, bias=True)
    (2): Linear(in_features=20, out_features=10, bias=True)
    (3): Linear(in_features=20, out_features=10, bias=True)
    (4): Linear(in_features=20, out_features=10, bias=True)
  )
)

当然，也可以使用 nn.ModuleList：

class multilayer(nn.Module):
    def __init__(self, num):
        super().__init__()
        layer_list = []
        for i in range(num):
            layer_list.append(nn.Linear(20, 10))
        self.ln = nn.ModuleList(layer_list)
    
    def forward(self, X):
        return self.ln(X)

multilayer(
  (ln): ModuleList(
    (0): Linear(in_features=20, out_features=10, bias=True)
    (1): Linear(in_features=20, out_features=10, bias=True)
    (2): Linear(in_features=20, out_features=10, bias=True)
    (3): Linear(in_features=20, out_features=10, bias=True)
    (4): Linear(in_features=20, out_features=10, bias=True)
  )
)

5.2 参数管理

有时我们希望提取参数，以便在其他环境中复用它们，将模型保存下来，以便它可以在其他软件中执行，或者为了获得科学的理解而进行检查。

本节，我们将介绍以下内容：

• 访问参数，用于调试、诊断和可视化；
• 参数初始化；
• 在不同模型组件间共享参数。

假定此时有一个单隐藏层的多层感知机

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)

tensor([[-0.5471], [-0.5554]], grad_fn=<AddmmBackward0>)

5.2.1 参数访问

同时，对于 Sequential 中，可以使用索引来访问模型的任意层，除此之外，可以使用 .state_dict() 来检查参数。比如，第二个全连接层的调用方法为 net[2].state_dict()。

OrderedDict([('weight', tensor([[-0.2183, -0.2935, -0.2471,  0.3105, -0.0285, -0.0140, -0.1047, -0.0894]])), ('bias', tensor([-0.0456]))])

1. 目标参数

parameter 是复合的类，包含值、梯度和额外信息。这就是我们需要显式参数值的原因。除了值之外，我们还可以访问每个参数的梯度。

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

<class 'torch.nn.parameter.Parameter'>
Parameter containing:
tensor([0.2615], requires_grad=True)
tensor([0.2615])

2. 一次性访问所有参数

当我们需要对所有参数执行操作时，逐个访问它们可能会很麻烦。当我们处理更复杂的块（例如，嵌套块）时，情况可能会变得特别复杂，因为我们需要递归整个树来提取每个子块的参数。下面，我们将通过演示来比较访问第一个全连接层的参数和访问所有层。

module.named_parameters 返回一个所有 module 参数的迭代器，返回参数名字和参数。

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

('weight', torch.Size([8, 4])) ('bias', torch.Size([8]))
('0.weight', torch.Size([8, 4])) ('0.bias', torch.Size([8])) ('2.weight', torch.Size([1, 8])) ('2.bias', torch.Size([1]))

也有另一种访问网络参数的方式：

net.state_dict()['2.bias'].data

tensor([0.2615])

3. 从嵌套块收集参数

def block1():
    return nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                        nn.Linear(8, 4), nn.ReLU())

def block2():
    net = nn.Sequential()
    for i in range(4):
        net.add_module(f'block {i}', block1())
    return net

rgnet = nn.Sequential(block2(), nn.Linear(4, 1))
rgnet(X)

tensor([[0.2608],
        [0.2611]], grad_fn=<AddmmBackward0>)

输出一下看看

print(rgnet)

Sequential(
  (0): Sequential(
    (block 0): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 1): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 2): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 3): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
  )
  (1): Linear(in_features=4, out_features=1, bias=True)
)

由于是嵌套了三层 Sequential 因此可以使用索引来访问层。

rgnet[0][1][0].bias.data

tensor([-0.0647,  0.1259, -0.3926, -0.3025, -0.1323,  0.3075,  0.4889,  0.1187])

5.2.2 参数初始化

深度学习框架提供默认随机初始化，也允许我们创建自定义初始化方法，满足我们通过其他规则实现初始化权重。

默认情况下，PyTorch 会根据一个范围均匀地初始化权重和偏置矩阵，这个范围是根据输入和输出维度计算出的。PyTorch 的 nn.init 模块提供了多种预置初始化方法。

1. 内置初始化

首先调用内置的初始化器。下面的代码将所有权重参数初始化为标准差为 \(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]

(tensor([-0.0261,  0.0005,  0.0169,  0.0050]), tensor(0.))

还可以将所有参数初始化为给定的常量，如初始化为 \(1\)。

def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 1)
        nn.init.zeros_(m.bias)
net.apply(init_constant)
net[0].weight.data[0], net[0].bias[0]

(tensor([1., 1., 1., 1.]), tensor(0., grad_fn=<SelectBackward0>))

我们还可以对某些块应用不同的初始化方法。例如，下面我们使用 Xavier 初始化方法初始化第一个神经网络层，然后将第三个神经网络层初始化为常量值 \(42\)。

def init_xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
def init_42(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 42)

net[0].apply(init_xavier)
net[2].apply(init_42)
print(net[0].weight.data[0])
print(net[2].weight.data)

tensor([ 0.3676,  0.3810,  0.5257, -0.0244])
tensor([[42., 42., 42., 42., 42., 42., 42., 42.]])

2. 自定义初始化

有时，深度学习框架没有提供我们需要的初始化方法。在下面的例子中，使用以下的分布为任意权重参数 \(w\) 定义初始化方法：

\[w \sim \begin{cases}
U(5, 10), &\text{可能性} \frac{1}{4} \\
0, &\text{可能性}\frac{1}{2} \\
U(-10, -5), &\text{可能性} \frac{1}{4}
\end{cases}
\]

同样，实现了一个 my_init 函数来应用到 net。

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]

Init weight torch.Size([8, 4])
Init weight torch.Size([1, 8])
tensor([[-7.2929, -0.0000, -0.0000, -5.2074],
        [ 9.1947, -8.8687,  0.0000,  0.0000]], grad_fn=<SliceBackward0>)

注意，始终可以直接设置参数。

net[0].weight.data[:] += 1
net[0].weight.data[0, 0] = 42
net[0].weight.data[0]

tensor([42.0000,  1.0000,  1.0000, -4.2074])

5.2.3 参数绑定

有时我们希望在多个层间共享参数：我们可以定义一个稠密层，然后使用它的参数来设置另一个层的参数。

# 我们需要给共享层一个名称，以便可以引用它的参数
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])
net[2].weight.data[0, 0] = 100
# 确保它们实际上是同一个对象，而不只是有相同的值
print(net[2].weight.data[0] == net[4].weight.data[0])

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

这个例子表明第三个和第五个神经网络层的参数是绑定的。它们不仅值相等，而且由相同的张量表示。因此，如果我们改变其中一个参数，另一个参数也会改变。这里有一个问题：当参数绑定时，梯度会发生什么情况？答案是由于模型参数包含梯度，因此在反向传播期间第二个隐藏层（即第三个神经网络层）和第三个隐藏层（即第五个神经网络层）的梯度会加在一起。

练习题

（1）使用之前没写的 NestMLP (FancyMLP) 模型访问各个层的参数。

class NestMLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(20, 64), nn.ReLU(),
                                 nn.Linear(64, 32), nn.ReLU())
        self.linear = nn.Linear(32, 16)

    def forward(self, X):
        return self.linear(self.net(X))

net = NestMLP()
for name, param in net.named_parameters():
    print(name, param.shape)

net.0.weight torch.Size([64, 20])
net.0.bias torch.Size([64])
net.2.weight torch.Size([32, 64])
net.2.bias torch.Size([32])
linear.weight torch.Size([16, 32])
linear.bias torch.Size([16])

（2）查看初始化模块文档以了解不同的初始化方法。

官方文档链接

（3）构建包含共享参数层的多层感知机并对其进行训练。在训练过程中，观察模型各层的参数和梯度。

举个简单的例子，\(z=wy, y=wx\)，不妨假设此时复制了两个与 \(w\) 相同的值 \(w_1, w_2\)。那么在反向传播中 \(\frac{\mathrm{d} z}{\mathrm{d} w} = \frac{\mathrm{d}z}{\mathrm{d} w_1} + \frac{\mathrm{d} z}{\mathrm{d} y} \frac{\mathrm{d} y}{\mathrm{d} w_2} = y + wx = 2wx\)，因此会是多倍梯度加和。

（4）为什么共享参数是个好方式？

可以减少参数，空间占用更小。但是正确性有待商榷。

5.3 延后初始化

延后初始化（defers initialization），即直到数据第一次通过模型传递时，框架才会动态地推断出每个层的大小。

在以后，当使用卷积神经网络时，由于输入维度（即图像的分辨率）将影响每个后续层的维数，有了该技术将更加方便。现在我们在编写代码时无须知道维度是什么就可以设置参数，这种能力可以大大简化定义和修改模型的任务。

延后初始化中只有第一层需要延迟初始化，但是框架仍是按顺序初始化的。等到知道了所有的参数形状，框架就可以初始化参数。

书上没有关于延后初始化的代码，原因在于 PyTorch 中的延后初始化层 nn.LazyLinear() 仍然还是一个开发中的 feature。所以这一节在 PyTorch 版的书里有什么存在的必要吗？

5.4 自定义层

本节将展示如何构建自定义层。

5.4.1 不带参数的层

首先，构造一个没有任何参数的自定义层。下面的 CenteredLayer 类要从其输入中减去均值。要构建它，我们只需继承基础层类并实现前向传播功能。

class CenteredLayer(nn.Module):
    def __init__(self):
        super().__init__()
    
    def forward(self, X):
        return X - X.mean()

5.4.2 带参数的层

下面继续定义具有参数的层， 这些参数可以通过训练进行调整。可以使用内置函数来创建参数，这些函数提供一些基本的管理功能。比如管理访问、初始化、共享、保存和加载模型参数。这样做的好处之一是：我们不需要为每个自定义层编写自定义的序列化程序。

下面实现自定义版本的全连接层：

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)

练习题

（1）设计一个接受输入并计算张量降维的层，它返回 \(y_k = \sum_{i,j} W_{ijk} x_i x_j\)

最好使用 transpose() 或者是 permute() 把 \(W_{ijk}\) 转换一个维度，变成 \(W_{kij}\)。这样就可以写成如下的形式了：

\[y_k = \boldsymbol{x}^T \boldsymbol{W}_k \boldsymbol{x}
\]

class testlayer1(nn.Module):
    def __init__(self, in_units, units):
        super().__init__()
        self.W = nn.Parameter(torch.randn(units, in_units, in_units))
    def forward(self, x):
        h1 = torch.matmul(x, self.W.data)
        h2 = torch.matmul(h1, x)
        return h2
    
net = testlayer1(4, 2)
a = torch.rand(4)
print(a, net(a))
# 验证一下第一个对不对
print(torch.matmul(a, torch.matmul(net.W[0], a)))

tensor([0.2971, 0.8508, 0.0615, 0.5073]) tensor([-0.5827, -1.1151])
tensor(-0.5827, grad_fn=<DotBackward0>)

第二题看不懂 QWQ

5.5 读写文件

5.5.1 加载和保存张量

本节内容为如何加载和存储权重向量和整个模型。

• torch.save(obj, f) 存储张量 obj 到 f 位置。
• torch.load(f) 读取 f 位置的文件。

书中给出了保存与读取张量、张量列表、张量字典的示例。

5.5.2 加载和保存模型参数

深度学习框架提供了内置函数来保存和加载整个网络。需要注意的一个重要细节是，这将保存模型的参数而不是保存整个模型。例如，如果有一个 \(3\) 层多层感知机，则需要单独指定架构。因为模型本身可以包含任意代码，所以模型本身难以序列化。因此，为了恢复模型，需要用代码生成架构，然后从磁盘加载参数。从多层感知机开始：

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)

接下来，将模型的参数 net.state_dict() 存储在一个 mlp.params 的文件中。

torch.save(net.state_dict(), 'mlp.params')

为了恢复模型，我们实例化了原始多层感知机模型的一个备份。这里不需要随机初始化模型参数，而是直接读取文件中存储的参数。

clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))

这样即完成了模型的保存和加载。

练习题

（1）即使不需要将经过训练的模型部署到不同的设备上，存储模型参数还有什么实际的好处？

可以让其他人复用模型，做重复实验。

（2）假设我们只想复用网络的一部分，以将其合并到不同的网络架构中。比如想在一个新的网络中使用之前网络的前两层，该怎么做？

这里仅使用上文中多层感知机的第一层作为例子。

old_net_state_dict = torch.load('mlp.params')
clone2 = MLP()
# 假设此处预处理剩下层已经完成
clone2.hidden.weight.data = old_net_state_dict["hidden.weight"]
clone2.hidden.bias.data = old_net_state_dict["hidden.bias"]

或者直接从这个基于 OrderedDict 的 state_dict 里面拿参数就行。

（3）如何同时保存网络架构和参数？需要对架构加上什么限制？

直接 torch.save(net) 即可。但是这个网络架构不包括 forward 函数。

5.6 GPU

可以使用 nvidia-smi 命令来查看显卡信息。

我用的 Kaggle 平台的 T4 2 张，可以完成本节的代码任务。

!nvidia-smi

Thu Apr 27 09:27:16 2023       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.161.03   Driver Version: 470.161.03   CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |
| N/A   36C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla T4            Off  | 00000000:00:05.0 Off |                    0 |
| N/A   34C    P8    10W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

5.6.1 计算设备

在 PyTorch 中，CPU 和 GPU 可以用 torch.device('cpu') 和 torch.device('cuda') 表示。应该注意的是，cpu 设备意味着所有物理 CPU 和内存，这意味着 PyTorch 的计算将尝试使用所有 CPU 核心。然而，gpu 设备只代表一个卡和相应的显存。如果有多个 GPU，我们使用 torch.device(f'cuda:{i}') 来表示第 \(i\) 块 GPU（\(i\) 从 \(0\) 开始）。另外，cuda:0 和 cuda 是等价的。

import torch
from torch import nn

torch.device('cpu'), torch.device('cuda'), torch.device('cuda:1')

(device(type='cpu'), device(type='cuda'), device(type='cuda', index=1))

还可以查询可用的 GPU 的数量。

torch.cuda.device_count()

2

原书中定义了两个方便的函数，这两个函数允许在不存在所需 GPU 的情况下运行代码。

• try_gpu(i) 尝试使用 \(i\) 号 GPU，如果存在返回 torch.device(f'cuda:{i}')，如果不存在返回 torch.device('cpu')。默认参数为 i=0。
• try_all_gpus() 尝试使用所有 GPU，如果存在 GPU 返回所有 GPU 的列表，如果不存在返回 [torch.device('cpu')]。

5.6.2 张量与 GPU

默认情况下，张量是在 CPU 上创建的。需要注意的是，无论何时我们要对多个项进行操作，它们都必须在同一个设备上。

1. 存储在 GPU 上

有几种方法可以在 GPU 上存储张量。例如，我们可以在创建张量时指定存储设备。接下来，我们在第一个 gpu 上创建张量变量 X。在 GPU 上创建的张量只消耗这个 GPU 的显存。我们可以使用 nvidia-smi 命令查看显存使用情况。 一般来说，我们需要确保不创建超过 GPU 显存限制的数据。

X = torch.ones(2, 3, device = try_gpu())
X

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')

假设还存在另一个 GPU，那么在另一个 GPU 上创建随机张量。

Y = torch.rand(2, 3, device = try_gpu(1))
Y

tensor([[0.4099, 0.3582, 0.8877],
        [0.7732, 0.8459, 0.1519]], device='cuda:1')

2. 复制

如果要计算 \(\sf X + Y\)，那么需要将它们弄到同一个设备上，然后才能执行运算操作。例如，下面的代码是将 \(\sf X\) 复制到第二个 GPU，然后执行加法运算。

Z = X.cuda(1)
print(X)
print(Z)

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')
tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:1')

当然，也可以使用 .to() 来执行复制：

Z = X.to(torch.device('cuda:1'))
Z

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:1')

相加：

Y + Z

tensor([[1.4099, 1.3582, 1.8877],
        [1.7732, 1.8459, 1.1519]], device='cuda:1')

假设变量 \(\sf Z\) 已经存在于第二个 GPU 上。如果我们还是调用 Z.cuda(1) 会发生什么？它将返回 \(\sf Z\)，而不会复制并分配新内存。

Z.cuda(1) is Z

True

注意调用 Z.to(torch.device("cuda:1")) is Z 也同样返回 True。

所以这个 .to() 和 .cuda() 有啥区别啊

5.6.3 神经网络与 GPU

类似地，可以神经网络模型可以指定设备。下面的代码将模型参数放在 GPU 上。

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())
net(X)

tensor([[-0.3980],
        [-0.3980]], device='cuda:0', grad_fn=<AddmmBackward0>)

练习题

只做第（4）题。

（4）测量同时在两个 GPU 上执行两个矩阵乘法与在一个 GPU 上按顺序执行两个矩阵乘法所需的时间。提示：应该看到近乎线性的缩放。

同时在两个 GPU 上执行矩阵乘法：

a = torch.rand(1000, 1000).to(try_gpu(0))
b = torch.rand(1000, 1000).to(try_gpu(0))
c = torch.rand(1000, 1000).to(try_gpu(1))
d = torch.rand(1000, 1000).to(try_gpu(1))
begintime = time.time()
for i in range(1000):
    e = torch.matmul(a, b)
    f = torch.matmul(c, d)
print(time.time() - begintime)

0.34023451805114746

在一个 GPU 上按顺序执行两个矩阵乘法所需的时间：

a = torch.rand(1000, 1000).to(try_gpu(0))
b = torch.rand(1000, 1000).to(try_gpu(0))
c = torch.rand(1000, 1000).to(try_gpu(0))
d = torch.rand(1000, 1000).to(try_gpu(0))
begintime = time.time()
for i in range(1000):
    e = torch.matmul(a, b)
    f = torch.matmul(c, d)
print(time.time() - begintime)

0.8642914295196533

差不多是两倍的差距。

原文链接：https://www.cnblogs.com/bringlu/p/17359969.html