nvmath-python (Beta) 是一个开源 Python 库，为 Python 程序员提供对 NVIDIA CUDA-X 数学库的高性能数学运算访问。nvmath-python 既提供底层库的低级绑定，也提供更高级别的 Python 抽象。它可与 PyTorch 和 CuPy 等现有 Python 软件包进行互操作。
在本文中，我将展示如何在 nvmath-python 中将 结语 与矩阵乘法结合使用。结语是可以与正在执行的数学运算（如 FFT 或矩阵乘法）融合的运算。可用的结语涵盖了大多数常见的深度学习计算。我通过实施简单神经网络的常见正向和反向传递运算来演示其用法。
要安装 nvmath-python，请 按照安装说明 操作。
使用 RELU_BIAS 后记优化正向传递
在本节中，我将演示如何使用 epilogs 实现简单线性层的前向传递。此层首先将输入向量乘以权重矩阵，然后向生成矩阵的每个元素添加偏差，最后应用 ReLU 激活函数。
ReLU 是修正线性单元的简称，是一种常用的激活函数，可以在保持正值不变的同时将负值替换为 0。
在矩阵运算方面，该层可以表示为：
在方程中，以下定义成立：
是一批形状为的输入向量
:
是层的输入数量。
是批量大小。
是形状的权重矩阵
:
是层的输出数量。
是长度为
假设您的输入、权重和偏差为 CuPy 数组：
num_inputs, num_outputs = 784, 100
batch_size = 256
weights = cupy.random.rand(num_outputs, num_inputs)
bias = cupy.random.rand(num_outputs)
x = cupy.zeros((num_inputs, batch_size))
在最基本的版本中，您可以通过使用 nvmath-python 计算 ，然后手动处理偏差和 ReLU 来实现此线性层，如下代码示例所示。
在本示例中，我使用 nvmath.linalg.advanced.Matmul -class.html” rel=”follow noopener” target=”_blank”>有状态 API ，其中您可以将初始化和规划与乘法的实际执行分开。当您必须执行多个类似的乘法运算时，我推荐这种方法，因为它可以让您分期偿还规划的初始成本。有关 Matmul 的更多信息，请参阅 nvmath.linalg.advanced.Matmul。
mm = Matmul(weights, x)
mm.plan()
def forward():
y = mm.execute()
y += bias[:,cupy.newaxis]
y[y < 0] = 0
return y
要提高代码的性能，请利用 RELU_BIAS epilog 在单个融合的 cuBLAS 操作中执行所有三个操作。这个结语首先将偏差添加到乘法结果中，然后应用 ReLU 函数。
您可以使用 Matmul.plan 方法的 epilog 参数指定结语。一些结语（包括 RELU_BIAS）会接收额外的输入，可在 epilog_inputs 字典中指定。有关结语的更多信息，请参阅 nvmath.linalg.advanced.Matmul 。
from nvmath.linalg.advanced import MatmulEpilog
mm = Matmul(weights, x)
mm.plan(epilog=MatmulEpilog.RELU_BIAS, epilog_inputs={"bias": bias})
def forward():
y = mm.execute()
return y
正如我稍后解释的那样，要通过 ReLU 函数进行反向传播，您必须知道向 ReLU 的哪些输入为正、哪些为负。此辅助信息称为 ReLU 掩码 ，可通过 RELU_AUX_BIAS 后记获得。
当使用带有辅助输出的结语时，Matmul.execute 将返回一个包含实际结果和辅助输出字典的元组。在 RELU_AUX_BIAS 的情况下，辅助输出字典只有一个键 relu_aux，其中包含 ReLu 掩码。该掩码是位编码的，可能难以读取，但在向后传递期间，有专门的结语可以为您执行此操作。
from nvmath.linalg.advanced import MatmulEpilog
mm = Matmul(weights, x)
mm.plan(epilog=MatmulEpilog.RELU_AUX_BIAS, epilog_inputs={"bias": bias})
relu_mask = None
def forward():
global relu_mask
y, aux_outputs = mm.execute()
relu_aux = aux_outputs["relu_aux"]
return y
使用 RELU_AUX_BIAS epilog 的实现速度比其朴素的实现要快，从而显著提升性能。
图 2 显示了对大小为（65536，16384）（16384，8192）的 float16 矩阵执行矩阵乘法运算，然后执行偏加和 ReLU 运算。性能在 NVIDIA H200 GPU 上进行测量。
使用 DRELU_BGRAD 后记优化反向传播
在神经网络的反向传播过程中，损失函数相对于输出的梯度会反向传播到网络层，以计算每个参数的梯度。
直观地说，对于每个操作，当其输出对损失的影响已知时，就有可能确定其输入和参数（例如权重矩阵中的值）如何影响损失。有关更多信息，请参阅 反向传播 。
在这一部分，我假设有多个线性层堆叠在一起。我对通常被认为属于不同层的操作序列实施反向传播：添加偏差、应用 ReLU 以及乘以权重。
forward 中实施的操作以及 backward 中涵盖的部分 让 
在反向传播中，当您知道 loss function 
，其中
，其中
是
，按批量维度求和
计算 Matmul，然后手动处理掩码和批量和来简单实现。
mm = Matmul(weights.T, grad)
mm.plan()
def backward():
grad_t1 = mm.execute()
grad_t1[mask] = 0 # assuming that `mask = (t1 < 0)`
grad_bias = cupy.sum(grad_t1, axis=1)
return grad_t1, grad_bias
要优化您的向后传递，请使用 DRELU_BGRAD 后记。假设梯度 grad 中可用。DRELU_BGRAD 的 epilog 需要一个输入 relu_aux，其中包含从 RELU_AUX_BIAS 的 epilog 返回的掩码。它将此遮罩应用于乘法结果。它还会返回一个辅助输出，其中包含结果的逐列总和，恰好是
mm = Matmul(weights.T, grad)
mm.plan(epilog=MatmulEpilog.DRELU_BGRAD, epilog_inputs={"relu_aux":relu_mask})
def backward():
grad_t1, aux_outputs = mm.execute()
grad_bias = aux_outputs["drelu_bgrad"]
return grad_t1, grad_bias
图 5 显示了对大小为（65536，16384）（16384，8192）的 float16 矩阵执行矩阵乘法运算，然后应用 ReLU 掩码和偏差梯度计算。该性能在 NVIDIA H200 GPU 上进行了测量。
结束语
借助 nvmath-python 的后记，您可以在 Python 代码中融合常见的深度学习计算，从而大幅提高性能。有关更多信息，请参阅 nvmath-python：在 Python 文档中充分发挥 NVIDIA Math Libraries 的功能 。
我们是一个开源库，请随时访问 /NVIDIA/nvmath-python GitHub 仓库并与我们联系。
