昇思MindSpore进阶教程--自动向量化Vmap(上)

大家好，我是刘明，明志科技创始人，华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享，如果你也喜欢我的文章，就点个关注吧

正文开始

AI融合计算的蓬勃发展，对框架能力提出了新的需求和挑战。问题场景和模型设计逐渐复杂化，使得业务数据的维度和计算逻辑的嵌套深度也相应增长。结合向量化优化手段可以有效优化性能瓶颈，但实现向量化优化对于普通用户并非易事。虽然用户可以很容易地实现低维数据运算逻辑，但随着数据维度的增长，业务逻辑也变得更为复杂，用户需要清晰了解各操作间的数据维度的逻辑映射关系，给用户的模型设计和编码带来了巨大挑战。自动向量化特性（Vmap）帮助用户解决了这个头疼的问题，该技术允许用户将特定的批处理逻辑从函数中剥离。用户在编写函数时，只需要先考虑低维的运算逻辑即可，通过调用vmap接口自动实现高维运算，并且支持嵌套调用，有效降低问题复杂度。

本教程介绍自动向量化vmap接口的使用方式，将模型或函数中高度重复的运算逻辑转换为并行的向量运算逻辑，从而获得更加精简的代码逻辑以及更高效的执行性能。

向量化思维

向量化思维在提升计算性能的技术中十分常见，向量化思维可公式化表示为：
$$\begin{array}{r}\begin{array}{c}{a}_1+b_1=c_1\\ a_2+b_2=c_2\\ a_3+b_3=c_3\\ a_4+b_4=c_4\end{array}\Rightarrow \vec{a}+\vec{b}=\vec{c}\end{array}$$
其核心思想在于将多次for循环的计算逻辑转换为一次对于向量的计算。如果将单个操作抽象为一个函数或者一个模型的操作集合，同样可应用向量化思维方式来处理。

手动向量化

首先，我们先构造一个简单的卷积函数，适用于一维向量场景：

import mindspore
from mindspore import Tensor, ops
import mindspore.numpy as mnp

x = mnp.arange(5).astype('float32')
w = mnp.array([1., 2., 3.])

def convolve(x, w):
    output = []
    for i in range(1, len(x) - 1):
        output.append(mnp.dot(x[i - 1 : i + 2], w))
    return mnp.stack(output)

convolve(x, w)

当我们期望该函数运用于计算一批一维的卷积运算时，我们很自然地会想到调用for循环进行批处理：

x_batch = mnp.stack([x, x, x])
w_batch = mnp.stack([w, w, w])

def manually_batch_conv(x_batch, w_batch):
    output = []
    for i in range(x_batch.shape[0]):
        output.append(convolve(x_batch[i], w_batch[i]))
    return mnp.stack(output)

manually_batch_conv(x_batch, w_batch)

很显然，通过这种实现方式我们能够得到正确的计算结果，但效率并不高。 当然，您也可以通过自己手动重写函数实现更高效率的向量化计算逻辑，但这将涉及对数据的索引、轴等信息的处理。

def manually_vectorization_conv(x_batch, w_batch):
    output = []
    for i in range(1, x_batch.shape[-1] - 1):
        output.append(mnp.sum(x_batch[:, i - 1 : i + 2] * w_batch, axis=1))
    return mnp.stack(output, axis=1)

manually_vectorization_conv(x_batch, w_batch)

在低维场景下，您可以很容易把握数据索引间的映射逻辑，但随着维度的增加，计算逻辑也变得更为复杂，您或许也会为此混乱的逻辑感到头疼。 幸运的是Vmap为我们提供了另一种实现方式。