【AI系统】动手实现 PyTorch 微分

简介: 本文介绍了使用操作符重载(OO)编程方式实现的自动微分,特别是采用反向模式(Reverse Mode)的实现方法。文中详细解释了操作符重载的基本概念及其在自动微分中的应用,以及反向模式的工作原理。通过 Python 示例代码,演示了如何手动实现类似 PyTorch 中自动微分的核心机制,包括定义 `Variable` 类、`Tape` 结构以及实现基本的数学运算符重载。最后,通过一个具体的数学函数示例展示了如何利用上述机制计算梯度,帮助读者理解反向模式自动微分的全过程。

这里记录一下使用操作符重载(OO)编程方式的自动微分,其中数学实现模式则是使用反向模式(Reverse Mode),综合起来就叫做反向 OO 实现 AD 啦。

基础知识

下面一起来回顾一下操作符重载和反向模式的一些基本概念,然后一起去尝试着用 Python 去实现 PyTorch 这个 AI 框架中最核心的自动微分机制是如何实现的。

操作符重载 OO

操作符重载:操作符重载或者称运算重载(Operator Overloading,OO),利用现代语言的多态特性(例如 C++/JAVA/Python 等高级语言),使用操作符重载对语言中基本运算表达式的微分规则进行封装。同样,重载后的操作符在运行时会记录所有的操作符和相应的组合关系,最后使用链式法则对上述基本表达式的微分结果进行组合完成自动微分。

在具有多态特性的现代编程语言中,运算符重载提供了实现自动微分的最直接方式,利用了编程语言的第一特性(first class feature),重新定义了微分基本操作语义的能力。

在 C++ 中使用运算符重载实现的流行工具是 ADOL-C(Walther 和 Griewank,2012)。ADOL-C 要求对变量使用启用 AD 的类型,并在 Tape 数据结构中记录变量的算术运算,随后可以在反向模式 AD 计算期间“回放”。Mxyzptlk 库 (Michelotti, 1990) 是 C++ 能够通过前向传播计算任意阶偏导数的另一个例子。

FADBAD++ 库(Bendtsen 和 Stauning,1996 年)使用模板和运算符重载为 C++ 实现自动微分。对于 Python 语言来说,autograd 提供正向和反向模式自动微分,支持高阶导数。在机器学习 ML 或者深度学习 DL 领域,目前 AI 框架中使用操作符重载 OO 的一个典型代表是 PyTorch,其中使用数据结构 Tape 来记录计算流程,在反向模式求解梯度的过程中进行 replay Operator。

下面总结一下操作符重载的一个基本流程:

  • 操作符重载:预定义了特定的数据结构,并对该数据结构重载了相应的基本运算操作符;
  • Tape 记录:程序在实际执行时会将相应表达式的操作类型和输入输出信息记录至特殊数据结构;
  • 遍历微分:得到特殊数据结构后,将对数据结构进行遍历并对其中记录的基本运算操作进行微分;
  • 链式组合:把结果通过链式法则进行组合,完成自动微分。

操作符重载法的优点可以总结如下:

  • 实现简单,只要求语言提供多态的特性能力;
  • 易用性高,重载操作符后跟使用原生语言的编程方式类似。

操作符重载法的缺点可以总结如下:

  • 需要显式的构造特殊数据结构和对特殊数据结构进行大量读写、遍历操作,这些额外数据结构和操作的引入不利于高阶微分的实现;
  • 对于类似 if,while 等控制流表达式,难以通过操作符重载进行微分规则定义。对于这些操作的处理会退化成基本表达式方法中特定函数封装的方式,难以使用语言原生的控制流表达式。

反向模式 Reverse Mode

反向自动微分同样是基于链式法则。仅需要一个前向过程和反向过程,就可以计算所有参数的导数或者梯度。因为需要结合前向和后向两个过程,因此反向自动微分会使用一个特殊的数据结构,来存储计算过程。

而这个特殊的数据结构例如 TensorFlow 或者 MindSpore,则是把所有的操作以一张图的方式存储下来,这张图可以是一个有向无环(DAG)的计算图;而 PyTorch 则是使用 Tape 来记录每一个操作,他们都表达了函数和变量的关系。

反向模式根据从后向前计算,依次得到对每个中间变量节点的偏导数,直到到达自变量节点处,这样就得到了每个输入的偏导数。在每个节点处,根据该节点的后续节点(前向传播中的后续节点)计算其导数值。

整个过程对应于多元复合函数求导时从最外层逐步向内侧求导。这样可以有效地把各个节点的梯度计算解耦开,每次只需要关注计算图中当前节点的梯度计算。

从下图可以看出来,reverse mode 和 forward mode 是一对相反过程,reverse mode 从最终结果开始求导,利用最终输出对每一个节点进行求导。下图虚线就是反向模式。

image

前向和后向两种模式的过程表达如下,表的左列浅色为前向计算函数值的过程,与前向计算时相同,右面列深色为反向计算导数值的过程。

反向模式的计算过程如图所示,其中:

$$ \overline{v_i}=\dfrac{\delta y}{\delta v_i}\\ $$

根据链式求导法则展开有:

$$ \frac{\partial f}{\partial x}=\sum_{k=1}^{N} \frac{\partial f}{\partial v_{k}} \frac{\partial v_{k}}{\partial \boldsymbol{x}}\\ $$

可以看出,左侧是源程序分解后得到的基本操作集合,而右侧则是每一个基本操作根据已知的求导规则和链式法则由下至上计算的求导结果。

反向操作符重载实现

下面的代码主要介绍反向模式自动微分的实现。目的是通过了解 PyTorch 的 auto diff 实现,来了解到上面复杂的反向操作符重载实现自动微分的原理,值的主要的是千万不要在乎这是 MindSpore 的实现还是 TensorFlow 版的实现(实际上都不是哈)。

首先,需要通过 typing 库导入一些辅助函数。

from typing import List, NamedTuple, Callable, Dict, Optional

_name = 1
def fresh_name():
    global _name
    name = f'v{_name}'
    _name += 1
    return name

fresh_name 用于打印跟 tape 相关的变量,并用 _name 来记录是第几个变量。

为了能够更好滴理解反向模式自动微分的实现,实现代码过程中不依赖 PyTorch 的 autograd。代码中添加了变量类 Variable 来跟踪计算梯度,并添加了梯度函数 grad() 来计算梯度。

对于标量损失 l 来说,程序中计算的每个张量 x 的值,都会计算值 dl/dX。反向模式从 dl/dl=1 开始,使用偏导数和链式规则向后传播导数,例如:

$$ dl/dx*dx/dy=dl/dy \\ $$

下面就是具体的实现过程,首先我们所有的操作都是通过 Python 进行操作符重载的,而操作符重载,通过 Variable 来封装跟踪计算的 Tensor。每个变量都有一个全局唯一的名称 fresh_name,因此可以在字典中跟踪该变量的梯度。为了便于理解,__init__ 有时会提供此名称作为参数。否则,每次都会生成一个新的临时值。

为了适配上面图中的简单计算,这里面只提供了乘、加、减、sin、log 五种计算方式。

class Variable:
    def __init__(self, value, name=None):
        self.value = value
        self.name = name or fresh_name()

    def __repr__(self):
        return repr(self.value)

    # We need to start with some tensors whose values were not computed
    # inside the autograd. This function constructs leaf nodes. 
    @staticmethod
    def constant(value, name=None):
        var = Variable(value, name)
        print(f'{var.name} = {value}')
        return var

    # Multiplication of a Variable, tracking gradients
    def __mul__(self, other):
        return ops_mul(self, other)

    def __add__(self, other):
        return ops_add(self, other)

    def __sub__(self, other):
        return ops_sub(self, other)

    def sin(self):
        return ops_sin(self)

    def log(self):
        return ops_log(self)

接下来需要跟踪 Variable 所有计算,以便向后应用链式规则。那么数据结构 Tape 有助于实现这一点。

class Tape(NamedTuple):
    inputs : List[str]
    outputs : List[str]
    # apply chain rule
    propagate : 'Callable[List[Variable], List[Variable]]'

输入 inputs 和输出 outputs 是原始计算的输入和输出变量的唯一名称。反向传播使用链式规则,将函数的输出梯度传播给输入。其输入为 dL/dOutputs,输出为 dL/dinput。Tape 只是一个记录所有计算的累积 List 列表。

下面提供了一种重置 Tape 的方法 reset_tape,方便运行多次自动微分,每次自动微分过程都会产生 Tape List。

gradient_tape : List[Tape] = []

# reset tape
def reset_tape():
    global _name
    _name = 1
    gradient_tape.clear()

现在来看看具体运算操作符是如何定义的,以乘法为例子啦,首先需要计算正向结果并创建一个新变量来表示,也就是 x = Variable(self.value * other.value)。然后定义了反向传播闭包 propagate,使用链规则来反向支撑梯度。

def ops_mul(self, other):
    # forward
    x = Variable(self.value * other.value)
    print(f'{x.name} = {self.name} * {other.name}')

    # backward
    def propagate(dl_doutputs):
        dl_dx, = dl_doutputs
        dx_dself = other # partial derivate of r = self*other
        dx_dother = self # partial derivate of r = self*other
        dl_dself = dl_dx * dx_dself
        dl_dother = dl_dx * dx_dother
        dl_dinputs = [dl_dself, dl_dother]
        return dl_dinputs

    # record the input and output of the op
    tape = Tape(inputs=[self.name, other.name], outputs=[x.name], propagate=propagate)
    gradient_tape.append(tape)
    return x
def ops_add(self, other):
    x = Variable(self.value + other.value)
    print(f'{x.name} = {self.name} + {other.name}')

    def propagate(dl_doutputs):
        dl_dx, = dl_doutputs
        dx_dself = Variable(1.)
        dx_dother = Variable(1.)
        dl_dself = dl_dx * dx_dself
        dl_dother = dl_dx * dx_dother
        return [dl_dself, dl_dother]

    # record the input and output of the op
    tape = Tape(inputs=[self.name, other.name], outputs=[x.name], propagate=propagate)
    gradient_tape.append(tape)
    return x

def ops_sub(self, other):
    x = Variable(self.value - other.value)
    print(f'{x.name} = {self.name} - {other.name}')

    def propagate(dl_doutputs):
        dl_dx, = dl_doutputs
        dx_dself = Variable(1.)
        dx_dother = Variable(-1.)
        dl_dself = dl_dx * dx_dself
        dl_dother = dl_dx * dx_dother
        return [dl_dself, dl_dother]

    # record the input and output of the op
    tape = Tape(inputs=[self.name, other.name], outputs=[x.name], propagate=propagate)
    gradient_tape.append(tape)
    return x

def ops_sin(self):
    x = Variable(np.sin(self.value))
    print(f'{x.name} = sin({self.name})')

    def propagate(dl_doutputs):
        dl_dx, = dl_doutputs
        dx_dself = Variable(np.cos(self.value))
        dl_dself = dl_dx * dx_dself
        return [dl_dself]

    # record the input and output of the op
    tape = Tape(inputs=[self.name], outputs=[x.name], propagate=propagate)
    gradient_tape.append(tape)
    return x

def ops_log(self):
    x = Variable(np.log(self.value))
    print(f'{x.name} = log({self.name})')

    def propagate(dl_doutputs):
        dl_dx, = dl_doutputs
        dx_dself = Variable(1 / self.value)
        dl_dself = dl_dx * dx_dself
        return [dl_dself]

    # record the input and output of the op
    tape = Tape(inputs=[self.name], outputs=[x.name], propagate=propagate)
    gradient_tape.append(tape)
    return x

grad 是将变量运算放在一起的梯度函数,函数的输入是 l 和对应的梯度结果 results。

def grad(l, results):
    dl_d = {
   } # map dL/dX for all values X
    dl_d[l.name] = Variable(1.)
    print("dl_d", dl_d)

    def gather_grad(entries):
        return [dl_d[entry] if entry in dl_d else None for entry in entries]

    for entry in reversed(gradient_tape):
        print(entry)
        dl_doutputs = gather_grad(entry.outputs)
        dl_dinputs = entry.propagate(dl_doutputs)

        for input, dl_dinput in zip(entry.inputs, dl_dinputs):
            if input not in dl_d:
                dl_d[input] = dl_dinput
            else:
                dl_d[input] += dl_dinput

    for name, value in dl_d.items():
        print(f'd{l.name}_d{name} = {value.name}')

    return gather_grad(result.name for result in results)

以下面公式为例:

$$ f(x_{1},x_{2})=ln(x_{1})+x_{1}*x_{2}−sin(x_{2}) \\$$

因为是基于操作符重载 OO 的方式进行计算,因此在初始化自变量 x 和 y 的值需要使用变量 Variable 来初始化,然后通过代码 f = Variable.log(x) + x * y - Variable.sin(y) 来实现。

reset_tape()

x = Variable.constant(2., name='v-1')
y = Variable.constant(5., name='v0')

f = Variable.log(x) + x * y - Variable.sin(y)
print(f)
    v-1 = 2.0
    v0 = 5.0
    v1 = log(v-1)
    v2 = v-1 * v0
    v3 = v1 + v2
    v4 = sin(v0)
    v5 = v3 - v4
    11.652071455223084

print(f) 可以看到是下面图中的左边正向运算,计算出前向的结果。下面的代码 grad(f, [x, y]) 就是利用前向最终的结果,通过 Tape 一个个反向的求解。得到最后的结果啦。

image

dx, dy = grad(f, [x, y])
print("dx", dx)
print("dy", dy)
    dl_d {'v5': 1.0}
    Tape(inputs=['v3', 'v4'], outputs=['v5'], propagate=<function ops_sub.<locals>.propagate at 0x7fd7a2c8c0d0>)
    v9 = v6 * v7
    v10 = v6 * v8
    Tape(inputs=['v0'], outputs=['v4'], propagate=<function ops_sin.<locals>.propagate at 0x7fd7a2c8c378>)
    v12 = v10 * v11
    Tape(inputs=['v1', 'v2'], outputs=['v3'], propagate=<function ops_add.<locals>.propagate at 0x7fd7a234e7b8>)
    v15 = v9 * v13
    v16 = v9 * v14
    Tape(inputs=['v-1', 'v0'], outputs=['v2'], propagate=<function ops_mul.<locals>.propagate at 0x7fd7a3982ae8>)
    v17 = v16 * v0
    v18 = v16 * v-1
    v19 = v12 + v18
    Tape(inputs=['v-1'], outputs=['v1'], propagate=<function ops_log.<locals>.propagate at 0x7fd7a3982c80>)
    v21 = v15 * v20
    v22 = v17 + v21
    dv5_dv5 = v6
    dv5_dv3 = v9
    dv5_dv4 = v10
    dv5_dv0 = v19
    dv5_dv1 = v15
    dv5_dv2 = v16
    dv5_dv-1 = v22
    dx 5.5
    dy 1.7163378145367738

如果您想了解更多AI知识,与AI专业人士交流,请立即访问昇腾社区官方网站https://www.hiascend.com/或者深入研读《AI系统:原理与架构》一书,这里汇聚了海量的AI学习资源和实践课程,为您的AI技术成长提供强劲动力。不仅如此,您还有机会投身于全国昇腾AI创新大赛和昇腾AI开发者创享日等盛事,发现AI世界的无限奥秘~

目录
相关文章
|
3天前
|
机器学习/深度学习 存储 人工智能
【AI系统】昇思MindSpore并行
本文介绍昇思MindSpore的并行训练技术,包括张量重排布、自动微分等,旨在简化并行策略搜索,提高大规模模型训练效率。文章探讨了大模型带来的挑战及现有框架的局限性,详细说明了MindSpore如何通过技术创新解决这些问题,实现高效的大模型训练。
46 20
【AI系统】昇思MindSpore并行
|
1天前
|
机器学习/深度学习 人工智能 自然语言处理
转载:【AI系统】AI的领域、场景与行业应用
本文概述了AI的历史、现状及发展趋势,探讨了AI在计算机视觉、自然语言处理、语音识别等领域的应用,以及在金融、医疗、教育、互联网等行业中的实践案例。随着技术进步,AI模型正从单一走向多样化,从小规模到大规模分布式训练,企业级AI系统设计面临更多挑战,同时也带来了新的研究与工程实践机遇。文中强调了AI基础设施的重要性,并鼓励读者深入了解AI系统的设计原则与研究方法,共同推动AI技术的发展。
转载:【AI系统】AI的领域、场景与行业应用
|
3天前
|
机器学习/深度学习 人工智能 分布式计算
【AI系统】混合并行
混合并行融合了数据并行、模型并行和流水线并行,旨在高效利用计算资源,尤其适合大规模深度学习模型训练。通过将模型和数据合理分配至多个设备,混合并行不仅提升了计算效率,还优化了内存使用,使得在有限的硬件条件下也能处理超大型模型。3D混合并行(DP+PP+TP)是最先进的形式,需至少8个GPU实现。此策略通过拓扑感知3D映射最大化计算效率,减少通信开销,是当前深度学习训练框架如Deepspeed和Colossal AI的核心技术之一。
42 15
【AI系统】混合并行
|
1天前
|
人工智能 安全 PyTorch
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
SPDL是Meta AI推出的开源高性能AI模型数据加载解决方案,基于多线程技术和异步事件循环,提供高吞吐量、低资源占用的数据加载功能,支持分布式系统和主流AI框架PyTorch。
23 10
SPDL:Meta AI 推出的开源高性能AI模型数据加载解决方案,兼容主流 AI 框架 PyTorch
|
3天前
|
存储 人工智能 PyTorch
【AI系统】张量并行
在大模型训练中,单个设备难以满足需求,模型并行技术应运而生。其中,张量并行(Tensor Parallelism, TP)将模型内部的参数和计算任务拆分到不同设备上,特别适用于大规模模型。本文介绍了张量并行的基本概念、实现方法及其在矩阵乘法、Transformer、Embedding和Cross Entropy Loss等场景中的应用,以及通过PyTorch DeviceMesh实现TP的具体步骤。
30 11
【AI系统】张量并行
|
3天前
|
存储 机器学习/深度学习 人工智能
【AI系统】完全分片数据并行 FSDP
本文深入探讨了AI框架中针对权重数据、优化器数据和梯度数据的分布式并行实现,特别是在PyTorch框架下的具体方案。文章首先回顾了通用数据并行和分布式数据并行的概念,重点讨论了同步与异步数据并行的差异。接着,文章详细介绍了如何在PyTorch中实现弹性数据并行,特别是完全分片数据并行(FSDP)的机制,包括其如何通过分片模型状态和剩余状态来减少内存消耗,提高训练效率。此外,文章还探讨了混合精度训练、损失缩放和内存消耗估算等关键技术,为理解和实施高效的分布式训练提供了全面的指导。
21 9
【AI系统】完全分片数据并行 FSDP
|
3天前
|
机器学习/深度学习 人工智能 PyTorch
【AI系统】数据并行
数据并行是一种在分布式AI系统中广泛应用的技术,通过将数据集划分成多个子集并在不同计算节点上并行处理,以提高计算效率和速度。在大规模机器学习和深度学习训练中,数据并行可以显著加快模型训练速度,减少训练时间,提升模型性能。每个计算节点接收完整的模型副本,但处理不同的数据子集,从而分摊计算任务,提高处理速度和效率。数据并行按同步方式可分为同步数据并行和异步数据并行,按实现方式包括数据并行、分布式数据并行、完全分片的数据并行等。其中,分布式数据并行(DDP)是当前应用最广泛的并行算法之一,通过高效的梯度聚合和参数同步机制,确保模型一致性,适用于大型NPU集群和AI系统。
39 7
【AI系统】数据并行
|
3天前
|
机器学习/深度学习 人工智能 前端开发
【AI系统】计算图的控制流实现
计算图作为有向无环图(DAG),能够抽象神经网络模型,但在编程中遇到控制流语句(如if、else、while、for)时,如何表示成为难题。引入控制流后,开发者可构建更复杂的模型结构,但部署含控制流的模型至不支持Python的设备上较为困难。目前,PyTorch仅支持Python控制流,而TensorFlow通过引入控制流原语来解决此问题。计算图的动态与静态实现各有优劣,动态图易于调试,静态图利于优化。
21 5
【AI系统】计算图的控制流实现
|
3天前
|
机器学习/深度学习 存储 人工智能
【AI系统】计算图与自动微分
自动求导利用链式法则计算雅可比矩阵,从结果节点逆向追溯计算路径,适用于神经网络训练中损失值对网络参数的梯度计算。AI框架中,自动微分与反向传播紧密相连,通过构建计算图实现高效梯度计算,支持动态和静态计算图两种模式。动态图如PyTorch,适合灵活调试;静态图如TensorFlow,利于性能优化。
25 6
【AI系统】计算图与自动微分
|
3天前
|
机器学习/深度学习 人工智能 算法
【AI系统】计算图挑战与未来
当前主流AI框架采用计算图抽象神经网络计算,以张量和算子为核心元素,有效表达模型计算逻辑。计算图不仅简化数据流动,支持内存优化和算子调度,还促进了自动微分功能的实现,区分静态图和动态图两种形式。未来,计算图将在图神经网络、大数据融合、推理部署及科学计算等领域持续演进,适应更复杂的计算需求。
32 5
【AI系统】计算图挑战与未来