第 5 章模型训练与复用 on 《从零学AI指南手册》

5.1 参数及日志管理

Mon, 01 Jan 0001 00:00:00 +0000

第 5 章模型训练与复用#

经过前面几个章节内容的介绍我们算是已经逐步迈入了深度学习的大门。所谓工欲善其事必先利其器，因此在接下来的这章内容中，我们将会逐一对深度学习模型训练过程中将会用到的一些辅助技能进行介绍，包括：如何有效对模型参数进行管理、怎么从本地文件中载入参数、如何保证模型训练过程的可追溯、模型的持久化与迁移方法以及模型的多GPU训练和预处理结果缓存等内容。

5.1 参数及日志管理#

在深度学习模型的实现过程中由于我们会频繁调整整个模型的超参数，例如需要突然新增一个丢弃率参数或者是模型的控制参数等等，而且这样的操作经常是跨多个函数或模块，如果依旧采用显示的参数名来传递参数就会变得十分复杂。如果模型参数数量较多，通过参数名来传递参数也会显得代码十分臃肿。

同时，由于在深度学习模型中通常会有较多的超参数，模型在训练过程中也会输出相应的评估结果、损失值甚至是部分权重参数结果等，为了使得整个模型的训练过程可追溯，因此就需要有效地将这些信息给保存下来以便不时之需。

5.1.1 参数传递#

例如对于某个深度学习模型来说，其训练部分的函数实现过程如下所示：

 1 def train(train_file_path=os.path.join('data', 'train.txt'),
 2           val_file_path=os.path.join('data', 'val.txt'),
 3           test_file_path=os.path.join('data', 'test.txt'),
 4           split_sep='_!_',
 5           is_sample_shuffle=True,
 6           batch_size=16,
 7           learning_rate=3.5e-5,
 8           max_sen_len=None,
 9           num_labels=3,
10           epochs=5):
11     dataset = get_dataset(train_file_path, val_file_path,max_sen_len,
12                           test_file_path, split_sep, is_sample_shuffle)
13     model = get_model(max_sen_len, num_labels)

从上述代码可以看出，第1~10行定义很多需要用到的参数，且在第11~13行中分别将这些参数传入到了对应的函数中。这样看起来似乎没有问题，但是此时你需要再添加一个参数到get_model()这个函数中，例如加入丢弃率来提高模型泛化能力，并且get_model()函数里面也是在不同模块都要用到丢弃率这个参数，如果直接采用新加参数的方式那难免会涉及到诸多地方的修改。

因此，对于模型参数有效管理的一种高效做法就是在所有地方均传入一个实例化的类对象，通过类对象访问类成员变量的方式来获取相应的参数值，这样在增删模型参数时只需要在原始类对象实例化的地方修改一次就能实现。首先，需要定义一个配置类，实现代码如下所示：

 1 class ModelConfig(object):
 2     def __init__(self,train_file_path=os.path.join('data', 'train.txt'),
 3                  val_file_path=os.path.join('data', 'val.txt'),
 4                  test_file_path=os.path.join('data', 'test.txt'),
 5                  split_sep='_!_',is_sample_shuffle=True,
 6                  batch_size=16,learning_rate=3.5e-5,
 7                  max_sen_len=None,num_labels=3,epochs=5):
 8         self.train_file_path = train_file_path
 9         self.val_file_path = val_file_path
10         self.test_file_path = test_file_path
11         self.split_sep = split_sep
12         self.is_sample_shuffle = is_sample_shuffle
13         self.batch_size = batch_size
14         self.learning_rate = learning_rate
15         self.max_sen_len = max_sen_len
16         self.num_labels = num_labels
17         self.epochs = epochs

在上述代码中便定义了模型所需要用到的参数，并且可以通过如下方式来进行访问，示例代码如下所示：

5.2 Tensorboard可视化

Mon, 01 Jan 0001 00:00:00 +0000

5.2 Tensorboard可视化#

在网络模型的训练过程中一般都需要通过观察模型损失值或准确率的变化趋势来确定模型的优化方向，例如学习率的动态调整、惩罚项系数等等。同时，对于图像处理方向来说可能还希望能够可视化模型的特征图或者是样本分类类别在空间中的分布情况等。虽然这些结果也可以在网络训练结果后取对应的变量通过matplotlib进行可视化，但是我们更希望在模型训练过程中就能对其各种状态进行可视化。

因此，对于上述需求我们可以借助谷歌开源的Tensorboard工具来进行实现。在接下来的这节内容中我们将会详细介绍如何在PyTorch中通过Tensorboard来对各类变量及指标进行可视化。

5.2.1 安装与启动#

如果需要在PyTorch中使用Tensorboard除了需要安装Tensorboard工具本身之外，还需要安装的便是TensorFlow本身。因为Tensorboard中的部分可视化功能在使用中会依赖于TensorFlow框架，例如 add_embedding()函数。

对于TensorFlow和Tensorboard的安装，我们只需要执行安装TensorFlow的命令便可以同时完成两者的安装：

1 pip install tensorflow

同时，由于只是借助于Tensorboard来进行可视化，因此在安装TensorFlow的时候不用区分是GPU还是CPU版本两者都可以，也就是说假如某台主机上装了GPU版本的PyTorch，而不管你是装的GPU版本还是CPU版的TensorFlow，Tensorboard都可以正常使用。

在安装成功之后通过可以通过如下命令来进行测试：

1 tensorboard --logdir=runs

并会出现如下提示：

1 TensorBoard 1.15.0 at http://localhost:6006/ (Press CTRL+C to quit)

此时便可以通过http://127.0.0.1:6006这个链接来访问Tensorboard的可视化页面，如图5-1所示。

图 5-1. Tensorboard启动成功界面图

如果发现打不开这个地址，可以尝试通过如下命令来进行启动，然后再通过http://127.0.0.1:6006这个链接来访问。

1 tensorboard --logdir=runs --host 0.0.0.0

其中--logdir用来指定可视化文件的目录地址，后续我们会详细介绍。

5.2.2 连接与访问#

上面我们介绍了如何在本地安装与启动Tensoboard，而更常见的一种场景便是在远程主机上运行代码但需要在本地电脑上查看可视化运行结果。如果需要实现这种目的通常来说有两种方法，下面我们分别进行介绍。

1. IP直接访问

在通过IP直接访问的方案中，不管是在类似于腾讯云或阿里云上租用的主机还是实验室的专用主机，在完成Tensoboard安装并启动后在自己电脑上都可以通过地址http://IP:6006来进行访问。但需要注意的是，上面的IP对于公网主机（如腾讯云）来说指的是主机的公网IP，对于实验室或学校的主机来说指的则是局域网的内网IP。同时，如果在远程主机上启动Tensoboard后发现在本地并不能够打开，那么可以通过如下方式来进行排查：

(1) 公网主机

在后台的安全策略里面查看一下6006这个端口有没有被打开，如果没有则需要打开；
查看IP是否为公网IP，在主机的后台管理页面可以看到。

(2) 局域网主机

查看本地电脑是否和主机处于同一网段；

查看主机的6006端口是否打开，如果没有可以参考如下命令打开

1 firewall-cmd --zone=public --list-ports # 查看已开放端口
2 firewall-cmd --zone=public --add-port=6006/tcp --permanent #开放6006端口
3 firewall-cmd --zone=public --remove-port=6006/tcp --permanent #关闭6006端口
4 firewall-cmd --reload   # 配置立即生效

2. 端口转发访问

当然，除了通过IP直接访问外还可以借助SSH反向隧道技术来进行访问。例如服务器只开了22端口而且你没有权限打开其它端口的情况。在这种情况下可以通过下面两种方式来进行远程连接：

(1) 命令行终端

如果你的命令行终端支持SSH命令（例如较新的Windows10 的CMD或者Linux等）的话，可以直接通过下面这一条命令来进行连接：

1 ssh -L 16006:127.0.0.1:6006 username@ip

这条命令的含义就是将服务器上6006端口的信息通过SSH转发到本地的16006端口，其中16006是本地的任意端口（无限制），只要不和本地应用有冲突就行，后面则是对应的用户名和IP。

5.3 模型的保存与复用

Mon, 01 Jan 0001 00:00:00 +0000

5.3 模型的保存与复用#

在深度学习中通常训练一个可用的模型都需要耗费极大的成本，因此在模型训练过程中就需要对满足某些条件下的网络权重参数进行保存，然后在实际推理过程中直接载入这些权重参数来完成模型的推理过程。同时，另外一种场景便是模型已经在一批数据上训练完成且完成了本地持久化保存，但可能过了一段时间后又收集到了一批新的数据，因此这时候就需要将之前的模型载入进行在新数据上进行增量训练或者是在整个数据上进行全量训练。

在PyTorch中可以通过torch.save()和torch.load()来完成上述场景中的主要步骤。下面，我们将以之前介绍的LeNet5网络模型为例来分别进行介绍。不过在这之前，我们先来看看PyTorch中模型参数的保存形式。

5.3.1 查看模型参数#

这里，我们依旧以4.4节内容中介绍的LeNet5网络模型为例进行讲解分析。在定义完LeNet5网络模型并完成实例化操作后，那么网络中对应的权重参数也都完成了初始化的工作，即有了一个初始值。同时，可以通过如下方式来访问：

1 import sys
2 sys.path.append("../../")
3 from Chapter04.C03_LeNet5.LeNet5 import LeNet5
4 if __name__ == '__main__':
5     model = LeNet5()
6     print("Model's state_dict:")
7     for (name, param) in model.state_dict().items():
8         print(name, param.size())

在上述代码中，第1~2行是将Chapter04这个搜索路径加入到系统路径中，否则第3行会提示” No module named ‘Chapter04’ “。第5行是实例化模型LeNet5，即初始化整个模型。第7~8行则是遍历模型中的每个参数。同时，需要注意的是通过model.state_dict()函数返回得到的是一个Python中的有序字段（OrderedDict），即是遍历输出的顺序结果就是元素插入字典时的顺序，例如这里插入的网络层。

上述代码运行结束后，其输出的结果为：

 1 Model's state_dict:
 2 conv.0.weight torch.Size([6, 1, 5, 5])
 3 conv.0.bias torch.Size([6])
 4 conv.3.weight torch.Size([16, 6, 5, 5])
 5 conv.3.bias torch.Size([16])
 6 fc.1.weight torch.Size([120, 400])
 7 fc.1.bias torch.Size([120])
 8 fc.3.weight torch.Size([84, 120])
 9 fc.3.bias torch.Size([84])
10 fc.5.weight torch.Size([10, 84])
11 fc.5.bias torch.Size([10])

在上述输出结果中，每一行的前半部分表示参数的名称，如conv.0.weight，后面部分表示该权重参数对应的形状。同时从输出结果可以看出，模型一共有5层权重参数，即conv.0、conv.3、fc.1、fc.3和fc.5。

5.3.2 自定义参数前缀#

在上面的输出结果中有两个地方值得注意：①参数名中的fc和conv前缀是根据定义LeNet5模型时nn.Sequential()时的名字所确定，即在第4.4.3节中定模型时使用了两个Sequential()实例对象，名称分别为conv和fc；②参数名中的数字表示每个Sequential()中网络层所在的位置。例如，如果将LeNet5网络结构定义成如下形式：

 1 class LeNet5(nn.Module):
 2     def __init__(self, ):
 3         super(LeNet5, self).__init__()
 4         self.lenet5 = nn.Sequential(  
 5             nn.Conv2d(1, 6, 5, padding=2),  
 6             nn.ReLU(), 
 7             nn.MaxPool2d(2, 2),  
 8             nn.Conv2d(6, 16, 5), 
 9             nn.ReLU(),
10             nn.MaxPool2d(2, 2),
11             nn.Flatten(),
12             nn.Linear(16 * 5 * 5, 120),
13             nn.ReLU(),
14             nn.Linear(120, 84),
15             nn.ReLU(),
16             nn.Linear(84, 10))

那么其参数名则为：

5.4 模型的迁移学习

Mon, 01 Jan 0001 00:00:00 +0000

5.4 模型的迁移学习#

在前面几节内容中，我们详细介绍了PyTorch中模型的保存及载入推理和复用等过程。在有了前期这些基础知识后，接下来我们再来介绍关于模型迁移学习（Transfer Learning）部分的内容。

5.4.1 迁移学习#

在深度神经网络中由于模型通常都含有大量的可学习参数，所以在训练数据不充分的情况下模型极易出现过拟合或者泛化能力差的情况。而另一方面，数据样本的标注又是一项既耗费时间又耗费财力的工作[6]，尤其是在一些需要业务专家介入的复杂任务标注中。因此，如何利用有限的数据来训练模型便成为了热门的研究方向。受到人类学习的启发——人类在学习并解决一个新问题的时候，总是可以依赖于先前所拥有的经验并迅速迁移到当前的场景中——研究人员开始提出一种两段式的学习框架，即先在一个通用的大规模数据集上训练得到一个预训练模型（Pre-trained Model），然后针对于特定的任务场景再根据少量的标注数据对整个模型进行微调（Fine-tuning），而这也被称为迁移学习。

在深度学习中迁移学习主要起源于图像处理领域，其背后的理念是如果一个模型是基于足够大且通用的数据集所训练得到，那么该模型将可以有效地充当视觉领域的通用模型，随后便可以直接将这些学习到的模型参数迁移到下游任务中而不必再从头开始训练整个模型[1]。

在图像处理领域中ImageNet是一个非常著名的大型通用数据集，它是由李飞飞团队于2007年所发起构建的一个项目，包含有超过1400万张手动标注的图片，旨在为世界各地的研究人员提供用于训练大规模物体识别模型的图像数据。[2]。自2010年以来，ImageNet项目每年都举办了一次大规模视觉识别挑战赛（ILSVRC），挑战赛使用了1000个类别的图片用于正确分类和检测目标及场景[3]。如图5-20所示便是ImageNet数据集中的部分图像。

图 5-20. ImageNet数据集示例图

根据4.2.3节内容可知，越是靠近输出层特征越抽象，越是靠近输入层的特征越具体。因此假如现在有一个开源的图片分类模型A是基于ImageNet数据集训练而来，如果在某任务场景中需要训练另外一个10分类模型B用于汽车型号的分类，那么便可以直接取模型A中的前若干层（靠近输入层）网络作为特征提取器，然后在此基础上再加入一个新的全连接分类层来构造得到模型B完成整个10分类任务，此时称模型A为预训练（Pre-trained）模型。同时，通常来说还可以根据是否让预训练模型中的参数参与整个模型的训练这两种方式来完成模型的迁移学习任务[4]。

在接下来的内容中，我们将会通过一个实际的示例来对模型的迁移学习过程进行介绍。以下完整示例代码可参见Code/Chapter05/C05_ModelTrans/文件夹。

5.4.2 模型定义与比较#

在4.4节内容中，我们详细介绍了LeNet5网络模型的原理及现实过程，并且同时根据5.3节的介绍我们也清楚了模型的保存与复用。现在假设有一个LeNet6网络模型，它是在LeNet5的基础上多增加了一个全连接层，此时便可以通过迁移学习将LeNet5模型中的部分参数用于LeNet6模型中。具体地，LeNet6模型结构的实现代码如下所示：

 1 class LeNet6(nn.Module):
 2     def __init__(self, ):
 3         super(LeNet6, self).__init__()
 4         self.conv = nn.Sequential(
 5             nn.Conv2d(in_channels=1, out_channels=6, 
 6                 kernel_size=5, padding=2),
 7             nn.ReLU(), nn.MaxPool2d(2, 2), nn.Conv2d(6, 16, 5), 
 8             nn.ReLU(),nn.MaxPool2d(2, 2))
 9         self.fc = nn.Sequential(
10             nn.Flatten(),nn.Linear(16 * 5 * 5, 120),
11             nn.ReLU(),nn.Linear(120, 84),nn.ReLU(),
12             nn.Linear(84, 64),nn.ReLU(),nn.Linear(64, 10))

在上述代码中，第1~11行是LeNet5模型的前4层。第12行便是LeNet6模型中新加入的一个网络层。

在模型定义结束后，便可以输出模型中对应的参数信息。同时，为了完成后续模型的迁移过程这里也将LeNet5保存在本地权重参数载入输出以便两者进行对比，示例代码如下所示：

 1 if __name__ == '__main__':
 2     print("\n=====Model paras in LeNet6:")
 3     model = LeNet6()
 4     for (name, param) in model.state_dict().items():
 5         print(name, param.size())
 6 
 7     model_save_path = os.path.join('../C04_ModelSaving', 'lenet5.pt')
 8     print("\n=====Model paras in LeNet5:")
 9     loaded_paras = torch.load(model_save_path)
10     for (name, param) in loaded_paras.items():
11         print(name, param.size())

在上述代码中，第2~5行是输出LeNet6模型中各个权重参数的名称和形状信息。第7~11行则是载入上一节中持久化保存到本地的LeNet5权重参数，并同时也输出每个参数的名称和形状。

5.5 开源模型复用

Mon, 01 Jan 0001 00:00:00 +0000

5.5 开源模型复用#

在前面两节内容中我们陆续介绍了在PyTorch框架中模型保存和迁移的基本原理，在接下来的这节内容中我们将以ResNet18在ImageNet上训练得到的1000分类预训练模型为例，将其迁移到CIFAR10数据集上进行微调。总体上来讲，我们首先需要实例化一个ResNet模型，然后再用预训练模型对其初始化；然后再将原始ResNet中最后一个1000分类的分类层改为CIFAR10数据对应10分类层；最后在CIFAR10数据集上完成整个模型的微调。以下完整示例代码可以参见Code/Chapter05/C06_PretrainedModel/文件。

5.5.1 ResNet结构介绍#

为了方便使用PyTorch官方开源的预训练模型，下面我们直接使用PyTorch框架中ResNet模型。同时，为了便于后续模型迁移理解，这里先简单介绍一下PyTorch中ResNet实现部分的代码。在PyTorch框架中，可以通过如下2代码来实例化一个残差网络，以ResNet18为例，示例代码如下：

1 from torchvision.models import resnet18
2 model = resnet18()

其中函数resnet18的实现过程为：

1 def resnet18(*, weights = None, ):
2     weights = ResNet18_Weights.verify(weights)
3     return _resnet(BasicBlock, [2, 2, 2, 2], weights,...)

在上述代码中，第2行是验证传入的预训练模型是否合法。第3行则是根据残差结构的数量返回ResNet18模型。

进一步，_restnet函数中的核心部分为：

1 model = ResNet(block, layers, **kwargs)

在上述代码中返回的便是一个残差网络的实例化对象，而类ResNet中的网络结构定义过程为：

 1 class ResNet(nn.Module):
 2     def __init__(self,......):
 3         super().__init__()
 4         ......
 5         self.layer1 = self._make_layer(block, 64, layers[0])
 6         self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
 7         self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
 8         self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
 9         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
10         self.fc = nn.Linear(512 * block.expansion, num_classes)

在上述代码中，第5~9行便是相应的残差结构和全局平均池化层。第10行则是对应最后的分类层，而有序将ResNet18迁移到CIFAR10数据集上需要修改的便是最后一个分类层。

5.5.2 迁移模型构造#

在清楚PyTorch中ResNet模型的基本实现结构之后，我们便可以对其进行修改以适应CIFAR10数据集。示例代码如下所示：

 1 from torchvision.models import resnet18
 2 from torchvision.models import ResNet18_Weights
 3 
 4 class ResNet18(nn.Module):
 5     def __init__(self, num_classes=10, frozen=False):
 6         super(ResNet18, self).__init__()
 7         self.resnet18 = resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
 8         if frozen:
 9             for (name, param) in self.resnet18.named_parameters():
10                 param.requires_grad = False
11                 logging.info(f"冻结参数: {name}, {param.shape}")
12         self.resnet18.fc = nn.Linear(512, num_classes)

在上述代码中，第7行是返回一个实例化的18层残差网络，同时指定了需要通过预训练模型来对其进行初始化。第8~11行则是用来判断是否需要对预训练部分的参数进行冻结，即不参与后续模型的训练过程，当然也可根据需要修改为对其中一部分参数进行冻结。第12行则是将原始残差网络的最后一层替换为符合新数据集的分类层。

5.6 多GPU训练

Mon, 01 Jan 0001 00:00:00 +0000

5.6 多GPU训练#

在深度学习中一些大型的网络模型往往需要大量的计算资源才能进行训练，因为每一层神经网络都需要对输入数据进行复杂的矩阵乘法和非线性变换操作。由于单个GPU的计算能力及显存有限可能无法满足大规模深度神经网络的训练需要，因此我们便需要使用多个GPU来加速网络的训练速度。在接下来的这节内容中，我们将会简单介绍几种多GPU模型训练的基本思想，并就其中一种最常见的方法进行详细讲解。

5.6.1 训练方式#

从理论上来讲，实现模型多GPU训练的策略有模型并行、数据并行和混合并行3种，然而在实际情况中并不是每一种都具有较高的可行性。

1. 模型并行

模型并行是指将模型的不同层或不同计算逻辑分配到不同GPU上进行训练的技术，此时每个 GPU 只负责处理部分逻辑的计算。通常这种方法适用于模型较大且无法在单个GPU上容纳的情况，例如现阶段各个大模型的训练过程通常都会采用模型并行技术。常见的模型并行技术有：层并行（Layer Parallelism），把模型的不同层分配给不同的 GPU；张量并行（Tensor Parallelism），把一个层内部的张量（如 Linear 权重矩阵）按维度切块分配到多个 GPU 上，如 LaMMA 和 DeepSeek 等就是采用的张量并行策略。

2. 数据并行

数据并行是指将输入网络的训练数据分成多个批次，每个批次在不同的GPU上进行并行计算。此时每个GPU上的模型权重都相同，只是处理的数据不同，每个GPU在训练完自己的批次数据后再将梯度更新汇总到主GPU上，从而实现模型参数的更新。这种方法的优点是简单易实现不容易出错，因此也是实现多GPU训练中使用最多的一种策略。

3. 混合并行

混合并行是一种同时使用数据并行和模型并行的技术。在混合并行中网络模型将会被拆分为多个子模型，并将每个子模型分配到不同的GPU上进行计算然后将计算好的结果传递给下一个GPU进行处理，同时在每个GPU也将使用数据并行技术进行处理。混合并行的优点在于它可以同时利用数据并行和模型并行的优势，因为数据并行可以处理大规模数据集，而模型并行可以扩展深度神经网络的规模。但混合并行也存在一些挑战，例如需要更多的硬件资源、实现难度较大、调试和优化复杂等问题。

以上便是3中并侧策略的基本思想，但是需要注意的是多 GPU 并不是越多越好，过多数量的 GPU 可能会造成通信延迟和资源浪费，并极有可能出现多个 GPU 的训练速度反而比单 GPU 更慢的情况。在实际使用中，我们需要根据具体的硬件条件和数据规模选择合适的多GPU训练策略。

下面，我们对使用最为常见的数据并行策略进行详细介绍。

5.6.2 数据并行#

在使用数据并行策略实现多GPU训练时，首先会将整个小批量数据再划分成多个小批次并分配到不同的GPU上，同时整个模型也将被复制到每个GPU上，然后在每个GPU上模型均各自独立地完成损失和梯度的计算，随后将每个GPU上计算得到的损失和梯度汇聚到主GPU上得到整个小批量数据样本的平均梯度，最后再将该梯度分配到其它GPU中进行各自模型参数的更新以完成一次迭代训练过程[1]。

图 5-23. 数据并行原理图

如图5-23所示便是含有两个GPU的数据并行原理图。例如此时每个小批量数据都含有256个样本，那么图示中每个GPU将会被分配得到128个样本进行后续的计算处理。同时，每个GPU上也都有着一模一样的网络模型，并且它们在各自拿到128个样本后会分别计算损失和梯度，然后再将两部分的梯度汇聚到主GPU上得到256个样本的平均梯度，最后再用该梯度通过梯度下降算法并行对每个GPU上的模型进行参数更新。

由此可以发现，对于数据并行这一多GPU训练策略来说，本质上就相当于每个GPU各自完成了部分数据样本的训练过程，且在整个前向传播和反向传播中每个GPU之间均是相互独立的，只有在进行整体损失和梯度的计算时才进行交互，因此基于数据并行的多GPU训练方法相对来说较为容易实现。但在实践中该方法也需要权衡计算资源、通信开销和同步效率等因素。

5.6.3 使用示例#

下面，我们以4.9节中介绍的ResNet18为例来介绍如何通过PyTorch框架实现网络模型的多GPU训练过程。在这里首先需要清楚的是，对于是否使用多GPU进行模型训练与模型的定义即前向传播过程无关，也就是我们只需要修改模型训练部分的代码即可。以下完整示例代码可参见Code/Chapter05/C07_MultiGPUs/train.py文件。

1. 获取GPU

首先，我们需要定义一个辅助函数来获取得到指定的GPU设备，示例代码如下所示：

1 def get_gpus(num=None):
2     gpu_nums = torch.cuda.device_count()
3     if isinstance(num, list):
4         devices = [torch.device(f'cuda:{i}')
5                    for i in num if i < gpu_nums]
6     else:
7         devices = [torch.device(f'cuda:{i}')
8                    for i in range(gpu_nums)][:num]
9     return devices if devices else [torch.device('cpu')]

在上述代码中，第1行num如果为list，则返回list中对应编号的GPU设备，如果num为整数，则返回主机中前num个GPU设备。第2行是得到当前主机上GPU设备的个数。第3~5行是根据num为list的情况获取对应的GPU设备。第7~8行是根据num为整数情况获取得到对应的GPU设备。第9行则是判断是否有GPU设备，没有则返回CPU设备。

5.7 数据预处理与缓存

Mon, 01 Jan 0001 00:00:00 +0000

5.7 数据预处理与缓存#

随着任务场景和深度学习模型的复杂化，使得模型在训练过程中每次调试时都需要花费较长的时间来等待数据集预处理结果。一个简单直接的办法就是在模型每次载入数据集时都预先判断本地是否有对应的缓存文件，如果有则直接载入，没有则重新处理并进行缓存。同时，为了方便这段处理逻辑能够方便地迁移到其它类似情况，因此我们需要将其定义成一个Python修饰器来进行使用。

下面，我们先来简单介绍一个Python中修饰器的功能及用法。

5.7.1 修饰器介绍#

关于什么是修饰器或装饰器（Decorator）我们这里就不从Python语法上来做详细的解释了。简单一句话，修饰器的作用的就是在正式执行某个功能函数之前，预先执行你想要执行的某些逻辑。例如在进行数据预处理之前先判断是否有对应的缓存文件。下面，我们直接从用法的层面来逐步了解Python中的修饰器。

首先来看这样一个场景，假如你之前已经定义了多个功能函数，但此时需要在日志文件中同时也输出每个函数的实际运行时间和其它相关信息。例如：

打印出当前主程序正在调用哪个功能函数的信息。例如：

1 def add(a=1, b=2):    
2     time.sleep(2)
3     r = a + b
4     return r
5 
6 def subtract(a=1, b=2):
7     time.sleep(3)
8     r = a - b
9     return r

在上述代码中，time.sleep(2)是为了模拟运行所花费的时间。

进一步，对于上述两个函数，如果需要实现打印运行时间等相关信息，可以通过如下类似方式实现：

1 def add(a=1, b=2):
2     print(f"正在执行函数 add() ！")
3     start_time = time.time()
4     time.sleep(2)
5     r = a + b
6     end_time = time.time()
7     print(f"一共耗时{(end_time - start_time):.3f}s")
8     return r

在上述代码中，第2、3、6和7行便是需要打印输出的相关信息。虽然通过这样的方式也能解决，但是如果有大量的函数都需要添加这么一段逻辑，那这种做法显然不可取。另外一种高效的方法则是使用Python中的修饰器。

假如现在已经定义好了一个名为get_info的修饰器，那么只需要通过如下方式便可以打印上述相关信息，示例代码如下所示：

 1 @get_info
 2 def subtract(a=1, b=2):
 3     time.sleep(3)
 4     r = a - b
 5     return r
 6 
 7 if __name__ == '__main__':
 8     subtract(3, 4)
 9 
10 正在执行函数 subtract() ！
11 一共耗时3.002s

在上述代码中，第1行便是调用了get_info修饰器。第2~5行是subtract函数原有的计算逻辑，并没有进行任何修改。所以，此时我们只需要在所有函数定义的地方使用get_info修饰器便可以实现运行时间计算的功能。

第 5 章 模型训练与复用 on 《从零学AI指南手册》

5.1 参数及日志管理

第 5 章 模型训练与复用#

5.1 参数及日志管理#

5.1.1 参数传递#

5.2 Tensorboard可视化

5.2 Tensorboard可视化#

5.2.1 安装与启动#

5.2.2 连接与访问#

5.3 模型的保存与复用

5.3 模型的保存与复用#

5.3.1 查看模型参数#

5.3.2 自定义参数前缀#

5.4 模型的迁移学习

5.4 模型的迁移学习#

5.4.1 迁移学习#

5.4.2 模型定义与比较#

5.5 开源模型复用

5.5 开源模型复用#

5.5.1 ResNet结构介绍#

5.5.2 迁移模型构造#

5.6 多GPU训练

5.6 多GPU训练#

5.6.1 训练方式#

5.6.2 数据并行#

5.6.3 使用示例#

5.7 数据预处理与缓存

5.7 数据预处理与缓存#

5.7.1 修饰器介绍#

第 5 章模型训练与复用 on 《从零学AI指南手册》

第 5 章模型训练与复用#