3 深度学习软件优化技术

by Posted on

注：本博客为本人（scutvk）纯原创内容，禁止私自转载。

随着深度学习应用的广泛传播，如何在软件层面上实现更高效的计算和降低内存也是模型轻量化需要考虑的问题。这些优化技术包括深度学习框架、编译器和库等方面的改进。我们将依次介绍这些方面的关键技术，并讨论它们的主要优缺点。

3.1 框架代码优化

此节包括一些优化深度学习框架代码的tricks。

3.1.1 算子融合

基本思想是将多个相邻的计算操作（算子）合并为一个更大的操作，从而减少计算中的冗余操作、提高计算效率和性能。

神经网络模型包含大量的计算操作，如加法、乘法、卷积等。这些操作通常是逐个执行的，但在某些情况下，将它们合并可以减少计算量、内存访问次数以及数据传输开销。

算子融合的优势包括：

  1. 提高计算性能：通过减少冗余操作和内存访问次数，算子融合可以降低计算时间。
  2. 减少内存占用：合并操作可以减少中间变量的使用，从而降低内存占用。
  3. 降低能耗：减少计算操作和内存访问可以降低整体能耗。

PyTorch中实现算子融合有两种方法：

  1. TorchScript：TorchScript是PyTorch的一个子模块，它可以将模型从Python代码转换为静态图表示。这使得模型可以在没有Python解释器的环境中运行，从而提高性能。TorchScript提供了一种方法，即将模型转换为ScriptModule，然后使用torch.jit.trace或torch.jit.script进行跟踪或脚本化。在这个过程中，TorchScript会自动进行一些优化，包括算子融合。这些优化有助于提高计算性能和减少内存占用。
  2. PyTorch FX（Functional Transformations）：PyTorch FX是另一个用于模型转换和优化的子模块。它可以捕获模型的计算图，并允许对计算图进行转换和优化。通过自定义优化策略，可以在PyTorch FX中实现算子融合。这使得可以更灵活地针对特定硬件和场景进行优化。

不过，由于硬件可以对默写计算有优化，算子融合也可能起到负优化的作用。

3.1.2 内存管理

神经网络模型通常具有数百万个参数，训练过程中也需要大量的中间值和梯度，需要大量的内存来存储这些参数及其相关信息。在目前最先进的GPU NVIDIA Tesla H100的显存为80G，但随着各种大模型的发展，模型参数和显存的增长明显不在一个数量级，由此，我们也需要关注内存管理技术。

内存交换

内存交换技术是指在加速设备内存和主存之间交换数据，通过在不使用变量时将其从加速设备的内存交换到主存的方式来降低加速设备的内存消耗，并在下一次访问变量之前将其交换回加速设备内存。

重计算

重计算技术的思想是将特征映射这样的中间结果在正向传播过程中及时地释放，在反向传播的计算需要用到特征映射时，再通过重新计算的方式生成，进而参与到当前计算中。

内存共享

内存共享技术指的是通过对不同变量生命周期的分析，在不同变量之间重复使用同一块内存空间。

参考：深度学习中的内存管理问题研究综述_唐名威的博客-CSDN博客

3.1.3 并行计算

由于大量的计算和数据处理需求，采用并行计算方法可以显著提高训练速度和效率。

并行加载

数据并行加载是指在深度学习模型训练时，将数据分成多个部分，并同时在多个CPU或计算设备上并行加载数据，以提高数据加载速度和效率。在PyTorch中，使用torch.utils.data.distributed.DistributedSampler模块可以实现数据并行加载。这个模块接收一个数据集作为输入，并将其划分成多个部分。每个部分都由一个进程负责加载和处理，提高数据加载速度和效率，同时也会减少每个进程的负载。

train_dataset = <加载数据集>
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) # 创建采样器
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size,
                                          shuffle=False, num_workers=num_workers,
                                          pin_memory=True, sampler=train_sampler) # 创建数据加载器
并行处理

在深度学习模型训练时，将数据分成多个部分，同时在多个GPU或计算设备上对这些数据进行并行处理，以提高训练速度和效率。在PyTorch中，使用torch.nn.DataParallel模块可以实现数据并行处理。这个模块接收一个模型作为输入，并自动将模型复制到多个GPU上进行训练。

3.2 编译器

在深度学习任务中，编译器扮演着至关重要的角色，它将高级语言编写的模型转换为针对特定硬件平台的低级指令。优化编译器能够提高计算效率，降低延迟，从而使得深度学习模型能够在各种硬件设备上实现更高性能。我们可以利用许多现成的、专门为深度学习任务设计的编译器。

  1. TVM

    TVM 是一个开源的深度学习编译器框架，旨在为各种硬件平台优化和部署深度学习模型。它起源于华盛顿大学的系统研究项目。TVM 提供了一个灵活、可扩展的编译架构，可以实现高效的模型优化和代码生成，从而在各种硬件设备上实现高性能深度学习推理。

    主要特点：

    1. 硬件无关：TVM 支持多种硬件平台，包括 CPU、GPU、FPGA 和专用 AI 芯片等。这使得开发者无需对特定硬件平台进行手动优化，能够自动适应不同的硬件设备。
    2. 多框架支持：TVM 支持多种深度学习框架，如 TensorFlow、PyTorch、MXNet、ONNX 等。开发者可以将这些框架训练出的模型导入到 TVM 中进行优化和部署。
    3. 端到端优化：TVM 采用了一种端到端的优化策略，从算子层面（如卷积、矩阵乘法等基本计算单元）到整个神经网络模型进行优化。这种优化策略可以有效地提高计算性能，降低内存占用和延迟。
    4. 自动调度：TVM 利用自动调度算法（AutoTVM 和 AutoScheduler）搜索最优的计算策略和数据布局，从而自动地为各种硬件平台生成高性能的代码。这大大降低了优化深度学习模型的门槛和成本。

  2. TC

    Tensor Comprehensions（TC）是一个开源的、用于高性能深度学习计算的编程框架。TC 旨在让开发者以一种更简洁、更直观的方式编写深度学习算子（例如卷积、矩阵乘法等）。TC 是由 Facebook AI Research (FAIR) 开发的，可以自动生成针对特定硬件的高效代码，尤其是针对 GPU 进行优化。

    主要特点：

    1. 基于 Polyhedral 模型：TC 使用了一种称为 Polyhedral 模型的数学表示方法，通过这种方法可以自动地为给定的计算任务找到高效的实现方式。Polyhedral 模型可以对循环嵌套结构进行深度分析，从而为优化提供强大的支持。
    2. 基于 Halide 语言：TC 使用了类似于 Halide 的领域特定语言（DSL），允许开发者用简洁的语法编写深度学习算子。通过这种语言，开发者可以轻松地描述计算过程和数据访问模式，而无需关心底层硬件实现的细节。
    3. 自动优化：TC 利用自动搜索算法，为给定的深度学习算子找到最优的映射策略，以实现高性能计算。这种自动优化方法降低了手动优化的复杂性和工作量。
  3. Glow

    Glow（Graph-Lowering Compiler）由Facebook AI Research（FAIR）开发，旨在为神经网络提供高效的硬件加速。Glow的主要目标是优化深度学习模型的性能，以便在各种硬件平台上实现最佳性能。Glow的设计灵活且可扩展，可以支持多种硬件设备。它的一个显著优势是其能够生成特定于硬件的优化代码，从而提高模型在不同硬件设备上的执行速度。

  4. nGraph+PlaidML

    nGraph 是一个用于深度学习框架的图编译器。它由英特尔公司开发，目的是为了优化各种深度学习框架（如 TensorFlow）在不同硬件平台上的性能。nGraph 的核心组件是一个高级的中间表示（IR），它可以表示计算图。nGraph 通过将框架的计算图转换为其中间表示，然后对其进行优化，以在目标硬件平台上实现更高的性能。

    PlaidML 是一个开源的设备无关的深度学习加速库，旨在为各种硬件平台提供统一的高性能计算支持。它支持许多不同的设备，如 NVIDIA、AMD 和英特尔的 GPU，以及其他像 CPU、FPGA 等硬件。PlaidML 的核心是一个高度优化的多后端计算引擎，它可以为各种设备生成高效的计算内核。

  5. XLA

    XLA（加速线性代数）是一个编译器和运行时库，旨在优化TensorFlow（后支持Pytorch、Keras）计算图。它由Google开发，用于提高TensorFlow在各种硬件平台上的性能。XLA通过将计算图转换为适合硬件特性的高效低级代码来实现这一目标。

参考：深度学习编译器整理 – 知乎 (zhihu.com)

3.3 计算库

  1. cuDNN

    cuDNN（CUDA Deep Neural Network library）是一个由NVIDIA开发的基于GPU的深度学习库，它提供了高度优化的底层函数，使得深度学习框架如TensorFlow、PyTorch、Caffe等能够充分利用NVIDIA GPU的计算能力。cuDNN极大地加速了深度学习训练和推理过程，使得开发者能够快速地构建和训练深度神经网络。

  2. BLAS

    BLAS（Basic Linear Algebra Subprograms）是一组线性代数基本操作的函数库。提供了一些高效的数学操作函数，例如向量和矩阵的乘法、矩阵的转置、向量的内积等，可以被广泛地用于各种科学计算、工程应用和数据分析等领域。

  3. cuBLAS

    cuBLAS（CUDA Basic Linear Algebra Subprograms）是一个NVIDIA开发的基于CUDA的线性代数库。它为GPU（图形处理器）上的线性代数计算提供了高性能的实现，主要针对NVIDIA GPU架构。cuBLAS库提供了BLAS规范中定义的各种线性代数操作，包括向量、矩阵以及复杂数的加法、减法、点乘、矩阵乘法等。cuBLAS的主要优势在于它充分利用了GPU的并行计算能力，使得线性代数运算能在GPU上高效执行。

  4. Intel-MKL

    Intel-MKL（Intel Math Kernel Library）是Intel 开发的一套数学库，旨在为高性能计算提供优化的数学函数。它针对英特尔处理器进行了高度优化，支持C、C++和Fortran编程语言，并提供了一套易于使用的API接口。

    Intel-MKL包含以下几个主要功能：

    1. 线性代数：Intel MKL提供了BLAS和LAPACK（线性代数包）库的优化实现。这些库包含用于向量、矩阵运算以及线性方程组求解的函数。
    2. 快速傅里叶变换（FFT）：Intel MKL提供了高效的FFT库，用于处理一维、多维和实数、复数数据的离散傅里叶变换。
    3. 向量数学：Intel MKL提供了一系列优化的向量数学函数，用于处理向量运算，如加法、减法、乘法、除法等。
    4. 稀疏矩阵运算：Intel MKL提供了一系列稀疏矩阵运算函数，如稀疏矩阵-向量乘法、稀疏矩阵-稠密矩阵乘法等。
    5. 非线性优化：Intel MKL提供了一套优化求解器，用于解决线性和非线性优化问题。

参考：

  1. 深度学习计算库(1)-概览BLAS、DNN加速库 – 知乎 (zhihu.com)
  2. 基本线性代数子程序 – 维基百科 (wikipedia.org)
  3. Math Kernel Library – Wikipedia

Published by vk

发表回复