28个PyTorch最佳实践技巧，全网最全！-Toy模板网

这篇具有很好参考价值的文章主要介绍了28个PyTorch最佳实践技巧，全网最全！。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

28个PyTorch最佳实践技巧，全网最全！

前言

本文为大家带来的是 28 个 Pytorch 的最佳实践技巧，在我力所能及的范围内尽可能收集的非常全了。大多数技巧我都用过了解过，但也有一些技巧我没用过也没了解过，比如和分布式优化相关的内容，像这部分内容大家可以选择性阅读，等到哪天用得上的时候知道在哪里能找到就可以。

废话不多说，下面开始正文👏👏👏

通用优化

首先给出的两条建议都是关于数据加载的。在Pytorch中，数据加载的核心程序是DataLoader，其构造函数的参数如下：

python复制代码DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

1.使用多进程数据加载

DataLoader 的默认设置是 num_workers = 0，这意味着数据加载是同步的，而且是在主进程中完成，也意味着主进程的训练必须等待数据可用才能继续执行，更意味着存在数据加载阻塞计算的可能。当然了，这种默认设置也有它的道理。即当用于在进程之间共享数据的资源（例如，共享内存、文件描述符）有限时，或者当整个数据集很小并且可以完全加载到内存中时，这种模式可能是首选。而且，这种模式下的错误跟踪更易读，因此对调试很有用。

为了避免数据加载阻塞计算代码，PyTorch 提供了一个简单的切换来执行多进程数据加载，只需设置 num_workers > 0 即可。在这种模式下，每次创建 DataLoader 的迭代器时（例如，当我们调用 enumerate(dataloader) 时），都会创建 num_workers 个工作进程。每个进程都接收以下初始化参数：dataset，collate_fn，和 worker_init_fn。

对于映射型（map-style）数据集，主进程使用采样器生成索引并将它们发送给工作进程。事实上，任何 shuffle 随机化都是在主进程中通过分配索引来指导加载完成的。

对于迭代型（iterable-style）数据集，由于每个工作进程都获得了 dataset 对象的副本，因此多进程加载通常会导致数据重复。用户可以使用 torch.utils.data.get_worker_info() 或者 worker_init_fn来单独配置每个副本。出于类似的原因，在多进程加载中，drop_last 参数会删除每个工作进程的迭代型数据集副本的最后一个非完整批次。

扩展：Pytorch 把数据集分成两类：map-style 和 iterable-style。区别是前者仅仅代表从索引/键到数据样本的映射，而后者是真正可迭代的数据样本。详情参见官方文档

多进程数据加载有两个好处：一是由于数据获取是异步的，因此在获取新数据批次时，主进程可以同时执行训练工作；二是由于多进程并行，加载大量数据的过程变得更快。但是，这种方式的缺点也很明显，就是容易爆内存，因此要根据自己的内存情况谨慎设置 num_workers 的值。

根据经验法则，num_workers 最好这样设置：num_workers = 4 * num_GPU。

2.使用固定内存

CUDA 把主机端的内存被分成两种，即可分页内存（pageable memroy）和页锁定内存（page-locked 或 pinned）。二者的区别如下：

内存类型	内存分配方式	分页	换页	DMA（Direct Memory Acess）访问
可分页内存	通过操作系统 API (malloc)	✔	✔	❌
页锁定内存	通过 CUDA API (cudaMallocHost/cudaHostAlloc)	❌	❌	✔

与可分页内存相比，页锁定内存的传输速率快，其一般是可分页内存的两倍左右。

默认情况下，主机 (CPU) 数据分配是可分页的。GPU 不能直接从可分页内存访问数据，因此当从可分页内存向设备内存传输数据时，CUDA 驱动程序必须首先分配一个固定内存，先把主机数据复制到固定内存中，再把数据从固定内存传输到设备内存中，如下图所示（Source）。

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

可以看出，固定内存仅仅用作设备内存与主机内存之间数据传输的暂存区。我们可以通过直接把数据存储在固定内存中来省略从可分页内存向固定内存拷贝数据这一步，从而加快速度。对于数据加载，只需设置 DataLoader 的参数 pin_memory = True 即可把获取的数据张量放入固定内存中，从而更快地将数据传输到支持 CUDA 的 GPU 中（异步内存复制）。

但是，默认的内存锁定逻辑仅识别张量以及包含张量的映射和迭代。因此，如果固定逻辑看到一个自定义类型的批处理（例如当 collate_fn 返回的是自定义批处理类型时），或者如果批处理的每个元素都是自定义类型，则固定逻辑将无法识别它们，它将返回该批次（或那些元素）而不固定内存。只有在自定义类型上定义 pin_memory() 方法，才能为自定义批处理或数据类型启用内存锁定。

示例如下：

python复制代码class SimpleCustomBatch:
    def __init__(self, data):
        transposed_data = list(zip(*data))
        self.inp = torch.stack(transposed_data[0], 0)
        self.tgt = torch.stack(transposed_data[1], 0)

    # 自定义类型上的自定义内存固定方法
    def pin_memory(self):
        self.inp = self.inp.pin_memory()
        self.tgt = self.tgt.pin_memory()
        return self

def collate_wrapper(batch):
    return SimpleCustomBatch(batch)

inps = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
tgts = torch.arange(10 * 5, dtype=torch.float32).view(10, 5)
dataset = TensorDataset(inps, tgts)

loader = DataLoader(dataset, batch_size=2, collate_fn=collate_wrapper,
                    pin_memory=True)

for batch_ndx, sample in enumerate(loader):
    print(sample.inp.is_pinned())
    print(sample.tgt.is_pinned())

关于使用多进程数据加载和固定内存带来的性能提升可以参见下图：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

固定内存本来应该放在 GPU 优化中，但是和数据加载相关就一块在这讲了。

3.禁用验证或推理时的梯度计算

PyTorch 从所有涉及需要梯度的张量的操作中保存中间缓冲区。通常，验证或推理不需要梯度。torch.no_grad() 上下文管理器可用于禁用指定代码块内的梯度计算，这可以加快执行速度并减少所需的内存量。torch.no_grad() 也可以用作函数装饰器。

示例：

python复制代码>>> x = torch.tensor([1], requires_grad=True)
>>> with torch.no_grad():
...   y = x * 2
>>> y.requires_grad
False
>>> @torch.no_grad()
... def doubler(x):
...     return x * 2
>>> z = doubler(x)
>>> z.requires_grad
False

4.禁用 batchnorm 前的卷积层中的偏差

torch.nn.Conv2d() 的参数偏差 bias 是默认为 True 的（ Conv1d 和 Conv3d 也是如此）。

如果在 nn.Conv2d 层之后直接跟的是 nn.BatchNorm2d 层，则不需要卷积中的偏差，也就是说nn.Conv2d 应该设置 bias = False。因为在第一步中 BatchNorm 减去了平均值，这有效地消除了偏差的影响。那么保留偏差就是多此一举，还不如将其禁用以减少模型的参数。

这也适用于 1d 和 3d 卷积，只要 BatchNorm 归一化（或其他归一化层）与卷积的偏差在相同的维度上。如果我们使用的是 torchvision 提供的模型，则没有这种顾虑，因为它已经帮我们实现了这种优化。

这种优化最适用于卷积层，因为它经常使用类似 Conv -> BatchNorm -> ReLU 这样的块。例如，下面是 MobileNet 使用的一个块：

python复制代码nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
nn.BatchNorm2d(hidden_dim),
nn.ReLU6(inplace=True),

而下面这个块则不适用此优化：

python复制代码nn.Linear(64, 32),
nn.ReLU(),
nn.BatchNorm2d(),

在这种情况下，尽管 nn.Linear 与 nn.BatchNorm2d 确实在相同的维度上有一个偏置项，但该项的影响被 ReLU 非线性压缩，因此没有实际的重复工作。

5.使用 None 将梯度归零，而不是 .zero_grad()

要将梯度归零，与其使用这种方法：

python复制代码model.zero_grad()
# 或
optimizer.zero_grad()

不如使用下面这种方法：

python复制代码for param in model.parameters():
    param.grad = None

第二种方法不会将每个单独参数的内存归零，随后的反向传递也使用赋值（写）而不是加法（读写）来存储梯度，这减少了内存操作的次数。

从 PyTorch 1.7 开始，我们也可以这么做：zero_grad(set_to_none=True)。这样，梯度被初始化为 None 而不是 0。

6.融合逐点操作

逐点运算（逐元素加法、乘法、数学函数 sin()、cos()、sigmoid() 等）可以融合到单个内核中，以摊销内存访问时间和内核启动时间。

PyTorch JIT 可以自动融合内核，尽管编译器中可能还有其他尚未实现的融合方法，并且并非所有设备类型都受到同等支持。

逐点操作是受内存限制的，对于每个操作 PyTorch 都会启动一个单独的内核。每个内核从内存中加载数据，执行计算（这一步通常代价很低）并将结果存储回内存中。

融合算子只为多个融合逐点操作启动一个内核，并且只将数据加载/存储一次到内存中。这使得 JIT 对于激活函数、优化器、自定义 RNN 单元等非常有用。

在最简单的情况下，我们可以通过将 torch.jit.script 装饰器应用于函数定义来启用融合，例如：

python复制代码@torch.jit.script
def fused_gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

7.为 CV 模型启用 channels_last 内存格式

PyTorch 1.5 引入了对卷积网络的 channels_last 内存格式的支持。这种格式旨在与 AMP 结合使用，以进一步加速使用 Tensor Cores 的卷积神经网络。

channels_last 内存格式是 NCHW 张量的另一种排序方式，这种格式仅适用于 4D NCWH 张量。

NCHW 张量的经典存储（例子中是两个具有 3 个颜色通道的 4x4 图像）如下图所示：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

采用 channels_last 内存格式则是下图这样：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

对 channels_last 的支持还处于 beta 阶段，但预计将适用于标准 CV 模型（例如 ResNet-50、SSD）。要将模型转换为 channels_last 格式，请参见 Channels Last Memory Format Tutorial。

8.检查点中间缓冲区

缓冲区检查点是一种减轻模型训练的内存容量负担的技术。它不是存储所有层的输入来计算反向传播中的上游梯度，而是存储几个层的输入，而其他层在反向传播期间重新计算。减少的内存需求可以增加批量大小，从而提高利用率。

应仔细选择检查点目标。最好不要存储重计算成本小的大层输出。示例目标层是激活函数（例如 ReLU、Sigmoid、Tanh）、上/下采样和具有小累积深度的矩阵向量操作。

PyTorch 支持原生的 torch.utils.checkpoint API 来自动执行检查点和重新计算。

9.禁用调试 API

许多 PyTorch API 用于调试，应在常规训练运行时禁用：

异常检测：torch.autograd.detect_anomaly ，torch.autograd.set_detect_anomaly(True)
剖析器相关：torch.autograd.profiler.emit_nvtx，torch.autograd.profiler.profile
自动求导梯度检查：torch.autograd.gradcheck，torch.autograd.gradgradcheck

10.每次梯度更新使用多个批次

在训练模式下通过神经网络前向传播一些值会创建一个计算图，为模型中的每个自由参数分配一个权重（梯度）。然后反向传播将这些梯度以调整为更准确更好的值，在此过程中消耗计算图。

然而，并不是每一次前向传递都需要一次反向传递。换句话说，在最终调用 loss.backward() 之前，我们可以随意调用 model(batch) 多次。计算图将继续累积梯度，直到您最终决定将它们全部折叠。

对于性能受到 GPU 内存瓶颈的模型，这种简单的技术提供了一种简单的方法来获得比内存更大的“虚拟”批量大小。

例如，如果 GPU 内存中每个批次只能容纳 16 个样本，则可以向前传递两个批次，然后向后传递一次，等价批量大小为 32。或者向前传递四次，向后传递一次等价批量大小为 64。

以下代码示例来说明其工作原理：

ini复制代码model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, labels)       # Compute loss function
    loss = loss / accumulation_steps                # Normalize our loss (if averaged)
    loss.backward()                                 # Backward pass
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()                           # Reset gradients tensors
        if (i+1) % evaluation_steps == 0:           # Evaluate the model when we...
            evaluate_model()                        # ...have no gradients accumulated

请注意，我们需要以某种方式组合每批损失，通常可以取平均值。

使用多批次梯度累计的缺点只有一个：就是训练期间发生的任何固定成本将会翻倍。例如，在主机和 GPU 内存之间传输数据时的延迟。

在 Lightning 中使用多批次梯度累计非常简单，只需要设置 accumulate_grad_batches 即可，例如：

scss复制代码trainer = Trainer(accumulate_grad_batches=16)
trainer.fit(model)

11.使用梯度剪裁

梯度裁剪是一种技术，最初是为处理 RNN 中的梯度爆炸而开发的，它将变得太大的梯度值裁剪到更现实的值。我们设置一个 max_grad，PyTorch 在反向传播时应用 min(max_grad, actual_grad)（注意：max_grad 为 10 将确保梯度值落在 [-10, 10] 范围内）。

梯度裁剪可以加快模型收敛，可以允许我们选择更高的学习率。感兴趣的可以参考这篇论文：Why Gradient Clipping Accelerates Training: A Theoretical Justification for Adaptivity

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

PyTorch 中的梯度剪裁是通过 torch.nn.utils.clip_grad_norm_ 提供的。您可以根据具体情况将其应用于各个参数组，但最简单和最常见的使用方法是将剪裁作为一个整体应用于模型：

less
复制代码torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad)

max_grad 设置为多大合适？没什么规律可循。最好将它视为一个超参数。2013 年的论文 Generating Sequences With Recurrent Neural Networks 在中间层该值设为 10，而在输出头该值设为 100。

12.使用 Lightning

PyTorch Lightning 是对 Pytorch 的抽象和包装。它的好处是可复用性强，易维护，逻辑清晰等。缺点也很明显，这个包需要学习和理解的内容还是挺多的。到底适不适合自己还是因人而异，大家不妨去试试。

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

13.使用 Profiler 分析代码

在 Lightning 当中，如果我们只想分析标准操作，我们可以设置 profiler="simple"。它使用内置的 SimpleProfiler。

ini复制代码# by passing a string
trainer = Trainer(..., profiler="simple")

# or by passing an instance
from pytorch_lightning.profiler import SimpleProfiler

profiler = SimpleProfiler()
trainer = Trainer(..., profiler=profiler)

分析器的结果将在训练 trainer.fit() 完成时打印。如下图所示：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

如果我们想了解有关每个事件期间调用的函数的更多信息，可以使用 AdvancedProfiler。此选项使用 Python 的 cProfiler 来提供在代码中调用的每个函数中花费的时间的深入报告。

ini复制代码# by passing a string
trainer = Trainer(..., profiler="advanced")

# or by passing an instance
from pytorch_lightning.profiler import AdvancedProfiler

profiler = AdvancedProfiler()
trainer = Trainer(..., profiler=profiler)

此报告可能很长，因此您还可以指定目录路径 dirpath 和文件名 filename 来保存报告，而不是将其记录到终端的输出中。下图显示了操作 get_train_batch 的分析：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

详情点击：链接

CPU 特定优化

14.利用非统一内存访问 (NUMA) 控件

NUMA 或非统一内存访问是一种用于数据中心机器的内存布局设计，旨在利用具有多个内存控制器和块的多插槽机器中的内存局部性。一般来说，所有深度学习工作负载（训练或推理）都可以获得更好的性能，而无需跨 NUMA 节点访问硬件资源。因此，推理可以在多个实例上运行，每个实例在一个套接字上运行，以提高吞吐量。对于单节点训练任务，推荐分布式训练，让每个训练过程运行在一个socket上。

在一般情况下，以下命令仅在第 N 个节点的内核上执行 PyTorch 脚本，并避免跨套接字内存访问以减少内存访问开销。

arduino
复制代码# numactl --cpunodebind=N --membind=N python <pytorch_script>

详情点击：链接

15.利用 OpenMP

OpenMP 用于为并行计算任务带来更好的性能。 OMP_NUM_THREADS 是可用于加速计算的最简单的开关。它确定用于 OpenMP 计算的线程数。 CPU 亲和性设置控制工作负载如何分布在多个内核上。它会影响通信开销、缓存行无效开销或页面抖动，因此正确设置 CPU 亲和性会带来性能优势。 GOMP_CPU_AFFINITY 或 KMP_AFFINITY 确定如何将 OpenMP* 线程绑定到物理处理单元。详情点击：链接。

使用以下命令，PyTorch 在 N 个 OpenMP 线程上运行任务。

shell
复制代码# export OMP_NUM_THREADS=N

通常，以下环境变量用于设置与 GNU OpenMP 实现的 CPU 亲和性。 OMP_PROC_BIND 指定线程是否可以在处理器之间移动。将其设置为 CLOSE 可使 OpenMP 线程保持靠近连续位置分区中的主线程。 OMP_SCHEDULE 确定 OpenMP 线程的调度方式。 GOMP_CPU_AFFINITY 将线程绑定到特定的 CPU。

shell复制代码# export OMP_SCHEDULE=STATIC
# export OMP_PROC_BIND=CLOSE
# export GOMP_CPU_AFFINITY="N-M"

16.英特尔 OpenMP 运行时库 (libiomp)

默认情况下，PyTorch 使用 GNU OpenMP (GNU libgomp) 进行并行计算。在英特尔平台上，英特尔 OpenMP 运行时库 (libiomp) 提供 OpenMP API 规范支持。与 libgomp 相比，它有时会带来更多的性能优势。利用环境变量 LD_PRELOAD 可以将 OpenMP 库切换到 libiomp：

shell
复制代码# export LD_PRELOAD=<path>/libiomp5.so:$LD_PRELOAD

与 GNU OpenMP 中的 CPU 关联设置类似，libiomp 中提供了环境变量来控制 CPU 关联设置。 KMP_AFFINITY 将 OpenMP 线程绑定到物理处理单元。 KMP_BLOCKTIME 设置线程在完成并行区域的执行之后，在休眠之前应该等待的时间（以毫秒为单位）。在大多数情况下，将 KMP_BLOCKTIME 设置为 1 或 0 会产生良好的性能。以下命令显示了英特尔 OpenMP 运行时库的常见设置。

shell复制代码# export KMP_AFFINITY=granularity=fine,compact,1,0
# export KMP_BLOCKTIME=1

17.切换内存分配器

对于深度学习工作负载，Jemalloc 或 TCMalloc 可以通过尽可能多地重用内存来获得比默认 malloc 函数更好的性能。Jemalloc 是一个通用的 malloc 实现，强调避免碎片和可扩展的并发支持。TCMalloc 还具有一些优化功能以加快程序执行速度。其中之一是在缓存中保存内存以加快对常用对象的访问。即使在解除分配之后持有此类缓存也有助于避免在以后重新分配此类内存时进行代价高昂的系统调用。使用环境变量 LD_PRELOAD 来利用其中之一。

shell
复制代码# export LD_PRELOAD=<jemalloc.so/tcmalloc.so>:$LD_PRELOAD

18.使用 DistributedDataParallel(DDP) 功能在 CPU 上训练模型

对于小规模模型或内存受限模型，例如 DLRM，在 CPU 上进行训练也是一个不错的选择。在具有多个套接字的机器上，分布式训练带来了高效的硬件资源使用来加速训练过程。Torch-ccl，使用英特尔® oneCCL（集体通信库）进行了优化，用于实现诸如 allreduce、allgather、alltoall 等集体的高效分布式深度学习训练，实现 PyTorch C10D ProcessGroup API，并且可以作为外部 ProcessGroup 动态加载。在 PyTorch DDP 模块中实现优化后，torhc-ccl 加速了通信操作。除了对通信内核进行的优化之外，torch-ccl 还具有同时计算通信功能。

GPU 特定优化

19.打开 cudNN 基准测试

这个优化技巧只适用于使用了大量卷积层的模型（例如，普通卷积神经网络，或具有 CNN 主干的模型架构）。

卷积层的核心是卷积运算，它是图像处理、信号处理、统计建模、压缩等应用的基础运算。所幸，我们已经开发了大量不同的算法可以在不同的阵列大小和硬件平台上有效地计算卷积。在 PyTorch 中，它基于 NVIDIA 的 cuDNN 框架来实现加速计算。cuDNN 有一个基准测试 API，它运行一个简短的程序来根据给定的输入大小和硬件选择最佳算法来执行卷积。

我们可以通过设置 torch.backends.cudnn.benchmark = True 来启用基准测试。开启后，第一次在我们的 GPU 设备上运行特定大小的卷积时，将首先运行快速基准程序以确定给定输入大小的最佳 cuDDN 卷积实现。此后，对相同大小矩阵的每个卷积运算都将使用该算法而不是默认算法。

关于使用 cudNN 基准测试带来的性能提升可以参见下图：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

注意：只有在保持输入大小（即传递给模型的批处理张量的形状）固定的情况下，使用 cudNN 基准测试才会提高速度。否则，每次输入大小更改时都会触发基准测试。好在绝大多数模型使用固定的张量形状和批量大小，因此这通常不成问题。

20.避免不必要的 CPU-GPU 同步

避免不必要的同步，尽可能让CPU跑在加速器前面，以保证加速器工作队列包含很多操作。

尽可能避免需要同步的操作，例如：

print(cuda_tensor)
cuda_tensor.item()
内存复制: tensor.cuda()，cuda_tensor.cpu()，tensor.to(device)
cuda_tensor.nonzero()
取决于在 cuda 张量上执行的操作的结果的 Python 控制流，例如：if (cuda_tensor != 0).all() 如果您尝试清除附加的计算图，请改用 .detach()。这不会将内存转移到 GPU，它会删除任何附加到该变量的计算图。

21.直接在目标设备上创建张量

不要像下面这样来生成随机张量：

scss
复制代码torch.rand(size).cuda()

因为，这首先会创建 CPU 张量，然后将其传输到 GPU，这真的很慢。我们应该直接在目标设备上生成输出：

ini
复制代码torch.rand(size, device=torch.device('cuda'))

如果我们使用的是 Lightning，它会自动把模型和批次放在正确的 GPU 上。但是，如果我们在代码中的某处创建一个新张量（例如：为 VAE 采样随机噪声或类似的东西），那么必须自己放置张量：

ini
复制代码torch.rand(size, device=self.device)

这适用于所有创建新张量并接受设备参数的函数：torch.rand()、torch.zeros()、torch.full() 等。

22.使用混合精度和 AMP

混合精度利用 Tensor Cores 并在 Volta 和更新的 GPU 架构上提供高达 3 倍的整体加速。要使用张量核心，应该启用 AMP，并且矩阵/张量维度应该满足调用使用张量核心的内核的要求。

要使用Tensor Cores：

将大小设置为 8 的倍数（映射到张量核心的维度）
启用 AMP 建议将精度调整为 16 bit，它有两个优点：

使用的内存减半，因为默认精度一般是 32 bit，这意味着我们可以将批量大小加倍并将训练时间缩短一半。
某些 GPU（V100、2080Ti）提供自动加速（快 3 到 8 倍），因为它们针对 16 位计算进行了优化。

在 Lightning 中，可以很容易启用 16 bit：

ini
复制代码Trainer(precision=16)

注意：在 PyTorch 1.6 之前，您还必须安装 Nvidia Apex。但是现在 16 位是 PyTorch 的原生版本。如果我们使用的是 Lightning，它两者都支持，能根据检测到的 PyTorch 版本自动切换。

23.在输入长度可变的情况下预分配内存

用于语音识别或 NLP 的模型通常在具有可变序列长度的输入张量上进行训练。对于 PyTorch 缓存分配器，可变长度可能会出现问题，并可能导致性能下降或意外的内存不足错误。如果一个序列长度较短的批次后面跟着另一个序列长度较长的批次，那么 PyTorch 将被迫从先前的迭代中释放中间缓冲区并重新分配新的缓冲区。此过程非常耗时，并且会导致缓存分配器中出现碎片，这可能会导致内存不足错误。

一个典型的解决方案是实现预分配。它由以下步骤组成：

生成具有最大序列长度的（通常是随机的）一批输入（对应于训练数据集中的最大长度或某个预定义的阈值）
使用生成的批次执行前向和后向传递，不执行优化器或学习率调度器，此步骤预分配最大大小的缓冲区，可在后续训练迭代中重复使用
将梯度归零
进行规则的训练

24.使用非阻塞设备内存传输

当我们向设备内存传输数据时，通过 .to(non_blocking=True) 启用异步（非阻塞）传输很有用，尤其是当使用带有 pin_memory=True 的 DataLoader 加载数据时。低级 CUDA 优化允许某些类型的数据从固定内存传输到 GPU 设备上，仅与 GPU 内核进程同时执行。其工作原理如下图(Source)：

28个PyTorch最佳实践技巧，全网最全！,pytorch,人工智能,python

在第一个（顺序）示例中，数据加载会阻止内核执行，反之亦然。在后两个（并发）示例中，加载和执行任务首先被分解为更小的子任务，然后以即时方式流水线化。

PyTorch 使用此功能将 GPU 代码执行与 GPU 数据传输一起流水线化：

ini复制代码# assuming the loader call uses pinned memory
# e.g. it was DataLoader(..., pin_memory=True)
for data, target in loader:
    # these two calls are also nonblocking
    data = data.to('cuda:0', non_blocking=True)
    target = target.to('cuda:0', non_blocking=True)
    # model(data) is blocking, so it's a synchronization point
    output = model(data)

在这个例子中，model(data) 是第一个同步点。PyTorch 和 CUDA 为我们创建数据和目标张量，将数据张量移动到 GPU，将目标张量移动到 GPU，然后在模型中执行前向传递。

分布式优化

25.使用 DistributedDataParallel 而不是 DataParallel

PyTorch 有两种方式来实现数据并行训练：

torch.nn.DataParallel
torch.nn.parallel.DistributedDataParallel

DistributedDataParallel 和 DataParallel 的区别在于：前者使用多进程，为每个 GPU 创建一个进程，而 DataParallel 使用多线程。通过使用多进程，每个 GPU 都有自己的专用进程，这避免了 Python 解释器的 GIL 带来的性能开销。More

26.如果使用 DistributedDataParallel 和梯度累积进行训练，则跳过不必要的 all-reduce

默认情况下，torch.nn.parallel.DistributedDataParallel 在每次反向传递后执行梯度全归约（all-reduce），以计算参与训练的所有工作人员的平均梯度。如果训练在 N 步上使用梯度累积，那么在每个训练步之后都不需要进行 all-reduce，只需要在最后一次调用 backward 之后，就在执行优化器之前执行 all-reduce。

DistributedDataParallel 提供 no_sync() 上下文管理器，它为特定迭代禁用梯度全归约。no_sync() 应用于梯度累积的前 N-1 次迭代，最后一次迭代应执行所需的梯度 all-reduce。

27.如果使用 DistributedDataParallel(find_unused_parameters=True)，则匹配构造函数中和执行期间的层顺序

带有 find_unused_parameters=True 的 torch.nn.parallel.DistributedDataParallel 使用模型构造函数中的层顺序和参数来构建用于 DistributedDataParallel 梯度全归约的存储桶。 DistributedDataParallel 将 all-reduce 与反向传递重叠。仅当给定存储桶中参数的所有梯度都可用时，才会异步触发特定存储桶的 All-reduce。

为了最大化重叠量，模型构造函数中的顺序应该与执行期间的顺序大致匹配。如果顺序不匹配，则整个bucket的all-reduce等待最后到达的梯度，这可能会减少backward pass和all-reduce之间的重叠，all-reduce最终可能会暴露出来，这减慢训练速度。

带有 find_unused_parameters=False（这是默认设置）的 DistributedDataParallel 依赖于基于反向传递期间遇到的操作顺序的自动存储桶形成。使用 find_unused_parameters=False 时，无需重新排序层或参数即可获得最佳性能。

28.分布式设置中的负载平衡工作负载

处理顺序数据的模型（语音识别、翻译、语言模型等）通常会发生负载不平衡。如果一个设备接收到一批数据的序列长度比其余设备的序列长度长，那么所有设备都会等待最后完成的工作人员。向后传递功能作为分布式设置与 DistributedDataParallel 后端的隐式同步点。

有多种方法可以解决负载均衡问题。核心思想是在每个全局批次中尽可能均匀地将工作负载分配给所有工作人员。例如，Transformer 通过形成具有大致恒定数量的令牌（以及批次中可变数量的序列）的批次来解决不平衡问题，其他模型通过对具有相似序列长度的样本进行分桶甚至通过按序列长度对数据集进行排序来解决不平衡问题。

正文到此结束，以下是 References：文章来源地址https://www.toymoban.com/news/detail-808122.html

Performance Tuning Guide — PyTorch Tutorials 1.10.1+cu102 documentation

7-tips-to-maximize-pytorch-performance

Tricks for training PyTorch models to convergence more quickly (spell.ml)

到了这里，关于28个PyTorch最佳实践技巧，全网最全！的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！