飞桨框架v2.3 API最新升级 | 对科学计算、概率分布和稀疏Tensor等提供更全面支持

本文已在飞桨公众号发布，查看请戳链接：
飞桨框架v2.3 API最新升级！对科学计算、概率分布和稀疏Tensor等提供更全面支持！

2022年5月飞桨框架2.3版本正式发布，相比飞桨框架2.2版本，API体系更加丰富，新增了100多个API，覆盖自动微分、概率分布、稀疏Tensor、拟牛顿优化器、线性代数、框架性能分析、硬件设备管理、视觉和语音领域等方面。整体上进一步丰富了飞桨框架动静统一、高低融合的API体系。

下面对每一类新增的API及其应用场景进行详细介绍，方便快速理解和上手。

新增自动微分API 以更好支持科学计算

科学计算许多计算场景需要嵌入物理信息，Loss往往通过高阶偏微分方程表示，如物理信息神经网络(Physics-Informed Neural Networks)。为了更好地支持科学计算场景，飞桨框架v2.3新增Jacobian、Hessian、jvp、vjp4个自动微分API，其中Jacobian、Hessian支持按行延迟计算，能够在复杂偏微分方程组中避免计算无用导数项，提升计算性能。目前已经在赛桨PaddleScience中得到应用。

基于延迟计算求解雅可比矩阵与海森矩阵

Jacobian(func, xs)计算func在 xs处的雅可比矩阵。为了降低计算整个矩阵的性能和显存开销，实现上采用了延迟计算方式，即按需计算部分行，并进行缓存。

import paddledef func(x, y):    return paddle.matmul(x, y)x = paddle.to_tensor([[1., 2.], [3., 4.]])J = paddle.incubate.autograd.Jacobian(func, [x, x])print(J[:, :])# Tensor(shape=[4, 8], dtype=float32, place=Place(gpu:0), stop_gradient=False,#        [[1., 3., 0., 0., 1., 0., 2., 0.],#         [2., 4., 0., 0., 0., 1., 0., 2.],#         [0., 0., 1., 3., 3., 0., 4., 0.],#         [0., 0., 2., 4., 0., 3., 0., 4.]])print(J[0, :])# Tensor(shape=[8], dtype=float32, place=Place(gpu:0), stop_gradient=False,#        [1., 3., 0., 0., 1., 0., 2., 0.])print(J[:, 0])# Tensor(shape=[4], dtype=float32, place=Place(gpu:0), stop_gradient=False,#        [1., 2., 0., 0.])

Jacobian延迟计算示例

Hessian(func, xs)计算func在xs的海森矩阵，func要求是vector-scalar函数，即输出是单个元素。海森矩阵是输出对输入的二阶导数，也可以理解为雅可比矩阵对输入的导数，因此实现上基于Jacobian，同样支持延迟计算与缓存。

def _jac_func(*xs):
    jac = Jacobian(func, xs, is_batched=is_batched)
    if (is_batched and jac.shape[1] != 1) or (not is_batched and
                                              jac.shape[0] != 1):
        raise RuntimeError(
            "The function given to Hessian shoud return as single element Tensor or batched single element Tensor."
        )
    return jac[:, 0, :] if is_batched else jac[0, :]

self.symbolic = Jacobian(_jac_func, xs, is_batched=is_batched)

Hessian实现核心代码

使用jvp和vjp计算前向微分与反向微分

jvp(func, xs, v）即雅可比矩阵-向量乘积，计算函数func在xs处的雅可比矩阵与向量v乘积。jvp主要用途是计算前向微分，前向微分即依据链式法则，从前向后遍历计算图，计算当前变量对输入的微分。

如图，输入x为n元向量，经过A、B、C得到标量输出y 利用jvp计算前向微分过程如下，其中J表示jvp计算。计算输出对所有输入的偏导需要n次前向计算过程。

vjp(func, xs, v）即向量-雅可比矩阵乘积，计算向量v与函数func在xs处的雅可比矩阵乘积。vjp主要用途是计算反向微分，反向微分即依据链式法则，从后向前遍历计算图，计算输出对当前变量的微分。

反向微分计算过程如下，其中J^T表示vjp计算过程。

体系化完善概率分布类API 以支持更多应用场景

概率分布在科学计算、强化学习、变分推断等场景中有着广泛应用。为了增强飞桨在概率编程方面能力，飞桨框架v2.3对概率分布基础设施及高层API进行统一规划与设计，提供低耦合、高内聚的基础设施以及更为灵活、完备的API功能。该类API总体数量扩充到25个，新增6个概率分布类、13个分布变换类及2个KL散度计算类API，并可持续扩展，丰富了随机采样算法，提供完备向量化语义支持，同时支持随机变量之间的变换，对于强化学习中策略梯度求解等场景，通过重参数化机制支持反向传播。

更为灵活完备的API体系

飞桨框架v2.3对概率分布API进行了整合与增强，新增指数族、狄利克雷等统计分布，提供13个概率分布变换以及KL散度计算功能，除此之外，基于基础分布和分布变换，提供构建高阶分布功能。

（1）概率统计与随机采样

新增Dirichelt、Beta、Multinomial三类分布，支持上述分布随机采样与概率统计，同时抽象ExponentialFamily分布基类，提供统一的熵计算、自然参数计算以及对数正则化等指数分布族相关基础功能。

import paddle 

beta = paddle.distribution.Beta(0.5, 0.5)
samples = beta.sample((10000,))

Beta 分布采样示例

（2）分布变换

分布变换类API用于将原分布的随机采样数据，经过线性/非线形变换生成目标分布。飞桨框架v2.3新增13个随机变量变换类API，提供随机变量之间可逆、可组合变换，如仿射变换AffineTransform、幂变换PowerTransform、指数变换ExpTransform等基础变换，以及变换之间的链式组合ChianTransform。

如下示例，将标准正态分布经过仿射变换，生成满足特定需求的目标分布。

import paddle
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns


normal = paddle.distribution.Normal(0., 1.)
samples = normal.sample([10000])

sns.distplot(samples)
plt.title('samples from a Norma(0., 1.)')
plt.show()

affine = paddle.distribution.AffineTransform(paddle.to_tensor(2.), paddle.to_tensor(2.))
transformed_samples = affine.forward(samples)
sns.distplot(transformed_samples)
plt.title('samples from a AffineTransform(2., 2.) Normal(0., 1.)')
plt.show()

（3）统一KL散度计算

提供统一的KL散度计算API，kl_divergence以及register_kl. kl_divergence用于计算相同或不同分布之间的KL散度，register_kl用于注册用户自定义的KL散度计算逻辑。

import paddle

p = paddle.distribution.Beta(alpha=0.5, beta=0.5)
q = paddle.distribution.Beta(alpha=0.3, beta=0.7)

print(paddle.distribution.kl_divergence(p, q))
# Tensor(shape=[1], dtype=float32, place=CUDAPlace(0), stop_gradient=True,
#        [0.21193528]

（4）构建高阶分布

除上述基础功能外，新增了TransformedDistribution(base: Distribution, transforms: Sequence[Transform])，用户可使用该API，基于一个基础分布和一系列分布变换，灵活构建更高阶分布。

import paddle 

# 基于标准正态分布和仿射变换，构建新的高阶分布
d = paddle.distribution.TransformedDistribution(
    paddle.distribution.Normal(0., 1.), 
    [paddle.distribution.AffineTransform(paddle.to_tensor(1.), paddle.to_tensor(2.))]
)

print(d.sample([10]))
# Tensor(shape=[10], dtype=float32, place=Place(gpu:0), stop_gradient=True,
#        [-0.10697651,  3.33609009, -0.86234951,  5.07457638,  0.75925219,
#         -4.17087793,  2.22579336, -0.93845034,  0.66054249,  1.50957513])
print(d.log_prob(paddle.to_tensor(0.5)))
# Tensor(shape=[1], dtype=float32, place=Place(gpu:0), stop_gradient=True,
#        [-1.64333570])

更为完善的向量化语义

传统使用概率分布采样，通常传入标量参数，构造一个分布实例。针对多元分布，通过多个独立一元分布构建；针对不同参数，构造不同参数的多个实例；针对批量采样，执行采样方法多次；最终将上述结果进行组合，得到想要结果。无法有效利用现代张量运算及对应GPU加速等能力。

通过如下三个维度将采样过程进行向量化：

• 采样形状(sample_shape) ：向量化批量采样过程；

• 批形状(batch_shape)：向量化不同参数的分布；

• 事件形状(event_shape)：向量化多元分布；

实现上，上述每个shape均为张量支持飞桨框架广播语义，最终样本数据形状为：

sample_data.shape = sample_shaple + batch_shape + event_shape

import paddle 

# 创建两个3元多项分布，分布1概率为[0.2, 0.3, 0.5]，分布2概率为[0.1, 0.9, 0.8]，因此事件形状为3，批形状为2
multinomial = paddle.distribution.Multinomial(total_count=10, probs=paddle.to_tensor([[0.2, 0.3, 0.5], [0.1, 0.9, 0.8]]))
print('event_shape: ', multinomial.event_shape)
print('batch_shape: ', multinomial.batch_shape)

# 采样形状为4，因此最终样本形状为[4, 2, 3]
samples = multinomial.sample((4,))
print('sample_shape: ', samples.shape)

# event_shape:  (3,)
# batch_shape:  (2,)
# sample_shape:  [4, 2, 3]

更为丰富的采样算法

随机采样是生成概率分布样本的关键过程。不同随机采样算法在应用场景、性能存在较大差距。飞桨框架v2.3在原有算法基础上，新增逆变换采样以及接受-拒绝采样算法，以支持不同概率分布采样需求。如多项分布使用逆变换采样比传统采样方法具备更高的接受率，狄利克雷分布由于其反函数解析式复杂，使用接受-拒绝采样能更好的模拟真实分布过程。除上述通用采样算法之外，在基础设施层，还保留了对特定领域优化采样算法扩展性，如对于正态分布，box-muller算法具备更好性能，可按需扩展到正态分布上。

支持反向传播

概率分布参数化目的让此类API接口具备反向传播能力，应用在策略梯度、变分推理等场景。因为随机采样方法sample不支持反向传播，考虑业界使用习惯，我们在Distribution基类中增加了重参数化采样rsample，将目标分布通过不含参数的标准分布的表示，采样过程可微。

如非标准正态分布使用标准正态分布表示，能够支持反向传播。

def rsample(self, sample_shape):    
    shape = self._extended_shape(sample_shape)    
    eps = _standard_normal(shape, dtype=self.loc.dtype, device=self.loc.device)    
    return self.loc + eps * self.scale

新增创建2种稀疏Tensor的API 支持和常规Tensor进行转换

通常的神经网络是稠密网络，但是随着模型大小的指数型爆炸，越来越多的科研工作者和公司都无法提供足够的资源去训练这些强大却昂贵的模型，而网络稀疏化已经被认证可以用来加快网络的训练，同时降低对硬件资源的要求。

目前越来越多的场景有稀疏训练的需求，比如NLP中的稀疏Attention，CV，3D点云等。对于深度学习框架，其涵盖了Sparse Tensor基础数据结构、Sparse Tensor 的计算类API/OP、Sparse Tensor的稀疏训练网络层三个方面的能力。飞桨框架v2.3已支持创建和使用2种稀疏Tensor。

COO稀疏格式

飞桨支持COO稀疏格式来存储Tensor，COO是最通用且简单的坐标存储方式，包含两个数据：非0元素坐标、非0元素值，如图示：

通过paddle.sparse.sparse_coo_tensor，指定 非0元素坐标(indices)、非0元素值(values) 可创建一个COO格式的稀疏Tensor，其支持2D及以上任意维度，同时还支持非0元素为DenseTensor的Hybird COO形式。其打印效果如下：

indices = [[0, 1, 1, 2, 3, 3], 
           [1, 0, 3, 2, 1, 3]]
values = [1, 2, 3, 4, 5, 6]
dense_shape = [4, 4]
coo = paddle.sparse.sparse_coo_tensor(indices, values, dense_shape)
print(coo)
# Tensor(shape=[4, 4], dtype=paddle.int64, place=Place(gpu:0), stop_gradient=True, 
#        indices=[[0, 1, 1, 2, 3, 3],
#                 [1, 0, 3, 2, 1, 3]], 
#        values=[1, 2, 3, 4, 5, 6])

CSR稀疏格式

飞桨支持CSR稀疏格式存储Tensor，CSR是压缩行信息存储格式，其记录了每一行中第一个非0元素的索引，按行进行压缩。这种格式在行切片、矩阵乘、矩阵与向量乘的场景中运算更为高效。

通过paddle.sparse.sparse_csr_tensor，指定每行首个非0元素索引(crows)、非0元素列信息(cols)、非0元素值(values) 可创建一个CSR格式的稀疏Tensor，其支持2D/3D维度。其打印效果如下：

crows = [0, 1, 3, 4, 6]
cols = [1, 0, 3, 2, 1, 3]
values = [1, 2, 3, 4, 5, 6]
dense_shape = [4, 4]
csr = paddle.sparse.sparse_csr_tensor(crows, cols, values, dense_shape)
print(csr)
# Tensor(shape=[4, 4], dtype=paddle.int64, place=Place(gpu:0), stop_gradient=True, 
#        crows=[0, 1, 3, 4, 6], 
#        cols=[1, 0, 3, 2, 1, 3], 
#        values=[1, 2, 3, 4, 5, 6])

稀疏与稠密Tensor互转

直接通过非0元素的索引信息来创建稀疏Tensor较为繁琐，同时在训练网络中可能同时存在DenseTensor与SparseTensor，飞桨支持通过Tesor.to_dense() 、Tesor.to_sparse_coo()、Tesor.to_sparse_csr()三个Tensor类成员API，实现Sparse与Dense之间交互。

dense = paddle.to_tensor([[0, 1, 0, 2], 
                          [0, 0, 3, 4]], dtype='float32')
coo = dense.to_sparse_coo(sparse_dim=2)
# Tensor(shape=[2, 4], dtype=paddle.float32, place=Place(gpu:0), stop_gradient=True, 
#       indices=[[0, 0, 1, 1],
#                [1, 3, 2, 3]], 
#       values=[1., 2., 3., 4.])
csr = dense.to_sparse_csr()
# Tensor(shape=[2, 4], dtype=paddle.float32, place=Place(gpu:0), stop_gradient=True, 
#       crows=[0, 2, 4], 
#       cols=[1, 3, 2, 3], 
#       values=[1., 2., 3., 4.]）

飞桨对稀疏Tensor的数据支持将为后续的Sparse Tensor 的计算类API/OP、Sparse Tensor的稀疏训练网络层两方面的能力奠定基础，未来飞桨将围绕这两个方面展开更多工作，并通过NLP中的稀疏Attention、3D点云模型验证，整体提升飞桨的稀疏场景应用能力。

新增拟牛顿法二阶优化API

拟牛顿法二阶优化器因为在参数优化中引入了更多信息，相比传统的一阶优化器(SGD、Adam等)在收敛速度和精度上有着明显优势。经典的二阶优化算法有BFGS（Broyden–Fletcher–Goldfarb–Shanno 算法）、L-BFGS（Limited-memory BFGS，限制存储的BFGS）和CG（Conjugate Gradient，共轭梯度）。

飞桨框架v2.3新增2个二阶优化器API：minimize_bfgs和minimize_lbfgs，二者均采用Strong-Wolfe线搜索方法寻找每次迭代的最佳步长。

这两个API采用函数式接口，在使用时，只需传入待优化的函数和起始坐标，即可求解出函数的最小值。例如下面的示例，使用BFGS找到了函数 func的最小值 (0, 0)。此外，二阶优化API还支持设置迭代次数、容差、初始逆Hessian矩阵、线搜索的初始步长等参数，具有极高的灵活度，方便有相关经验的用户进行算法调优。

import paddle

def func(x):
    return paddle.dot(x, x)

x0 = paddle.to_tensor([1.3, 2.7])
results = paddle.incubate.optimizer.functional.minimize_bfgs(func, x0)
print("is_converge: ", results[0])
print("the minimum of func is: ", results[2])
# is_converge:  is_converge:  Tensor(shape=[1], dtype=bool, place=Place(gpu:0), stop_gradient=True,
#        [True])
# the minimum of func is:  Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True,
#        [0., 0.])

目前，二阶优化API已经应用在赛桨PaddleScience圆柱绕流模型训练中，并取得了不错的效果。从下方的loss-step曲线可以看出，在同样的迭代次数，使用二阶优化器比一阶优化器能使得Loss下降更快。

其他新增API 扩展线性代数API

为了更好地支持科学计算相关需求，飞桨框架v2.3扩展了线性代数的API，包括：计算线性方程组的最小二乘解paddle.linalg.lstsq，计算向量间的协方差 paddle.linalg.cov，通过Cholesky分解来计算具有唯一解的线性方程组 paddle.linalg.cholesky_solve，计算矩阵 lu 分解 paddle.linalg.lu以及对lu分解的结果进行解压缩paddle.linalg.lu_unpack。

（1）框架性能分析API

为了方便查找和使用，性能分析的9个API均整合到paddle.profiler.*，提供对训推过程中性能数据的收集，导出和统计的功能。支持自定义打点，支持将性能数据保存为 google chrome tracing 格式的文件，可在chrome中使用tracing插件查看。

（2）硬件设备管理API

paddle.device.cuda下新增：

• max_memory_allocated

• max_memory_reserved

• memory_allocated和memory_reserved

能够实时查看和分析模型显存占用。

新增get_device_name和get_device_capability：能够获取GPU设备名称信息和计算能力的主要和次要修订号。

（3）视觉和语音领域高层API

基于飞桨的基础API，大量扩充了paddle.vision.models.*常用视觉模型，用户可直接使用 AlexNet、DenseNet、GoogLeNet、InceptionV3、MobileNetV3、ResNeXt、ShuffleNetV2、SqueezeNet、WideResNet 等模型。

新增Viterbi 解码API：

• paddle.text.ViterbiDecoder和paddle.text.viterbi_decode

已用于序列标注模型的预测。

除了上面的介绍外，飞桨框架v2.3还扩充了大量Tensor计算相关的API，更好地支持业界论文中模型的实现，加速创新。详细列表可参考v2.3 release note，请点击下方链接获取。

v2.3 release note

结语

飞桨框架 2.3 版本提供了更加丰富的API体系，不仅更好地支撑深度学习图像、语言、语音领域的快速迭代和创新，而且在不断扩展对科学计算、强化学习、3D点云等场景应用的支持，同时也不断优化飞桨API的用户使用体验，使得API更易记、易查和易用，让深度学习技术的创新与应用更简单！

关注【飞桨PaddlePaddle】公众号
获取更多技术内容~