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

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

文章上半部分请查看
自动向量化Vmap(上)

自动向量化

Vmap可以帮助我们隐藏批处理维度，您只需要调用一个接口便可以将函数转换为向量化形式。

from mindspore import vmap

auto_vectorization_conv = vmap(convolve)
auto_vectorization_conv(x_batch, w_batch)

Vmap除了为您提供了简易的编程体验外，将循环逻辑下沉至函数的各个基元操作中，结合分布式并行优化以获得更高的执行性能。 默认情况下，vmap的输入输出沿第一个轴进行批处理，如果您的输入和输出并不总是期望沿0轴批处理，可以通过in_axes和out_axes参数进行指定。您可以为所有输入或输出位置分别指定批处理轴索引，也可以为所有输入或输出指定相同的批处理轴索引。

w_batch_t = ops.transpose(w_batch, (1, 0))

auto_vectorization_conv = vmap(convolve, in_axes=(0, 1), out_axes=1)
output = auto_vectorization_conv(x_batch, w_batch_t)

ops.transpose(output, (1, 0))

对于多个输入的场景，您还可以指定只对其中的某些入参进行批处理，如上述场景变为求一组一维向量与某一权重的卷积，可在in_axes参数中的输入对应位置配置None即可，None表示不沿任何轴进行批处理。

auto_vectorization_conv = vmap(convolve, in_axes=(0, None), out_axes=0)
auto_vectorization_conv(x_batch, w)

高阶函数的嵌套

Vmap本质上是一种高阶函数，它将函数作为输入，并返回可应用于批处理数据的向量化函数。用法上它允许和其他框架提供的高阶函数进行嵌套组合使用。

• vmap与vmap嵌套使用，应用于两层以上的批处理逻辑。

hyper_x = Tensor([[1., 2., 3., 4., 5.], [2., 3., 4., 5., 6.], [3., 4., 5., 6., 7.]], mindspore.float32)
hyper_w = Tensor([[1., 1., 1.], [2., 2., 2.], [3., 3., 3.]], mindspore.float32)

hyper_vmap_ger = vmap(vmap(convolve, in_axes=[None, 0]), in_axes=[0, None])
hyper_vmap_ger(hyper_x, hyper_w)

• grad内部嵌套vmap使用，应用于计算向量化函数的梯度等场景。

from mindspore import grad

def forward_fn(x, y):
    out = x + 2 * y
    out = ops.sin(out)
    reduce_sum = ops.ReduceSum()
    return reduce_sum(out)

x_hat = Tensor([[1., 2., 3.], [2., 3., 4.]], mindspore.float32)
y_hat = Tensor([[2., 3., 4.], [3., 4., 5.]], mindspore.float32)

grad_vmap_ger = grad(vmap(forward_fn), grad_position=(0, 1))
grad_vmap_ger(x_hat, y_hat)

• vmap内部嵌套grad使用，应用于计算批量梯度、高阶梯度计算等场景，如计算Jacobian矩阵。

vmap_grad_ger = vmap(grad(forward_fn, grad_position=(0, 1)))
vmap_grad_ger(x_hat, y_hat)

本教程中只简单介绍两层高阶函数组合嵌套的用法，您可以根据场景需求实现更多层次的嵌套。

Cell的自动向量化

之前的用例我们都是以函数对象作为输入，下面将介绍Cell对象结合vmap的用法。这是一个简单定义的全连接层的例子。

import mindspore.nn as nn
from mindspore import Parameter
from mindspore.common.initializer import initializer

class Dense(nn.Cell):
    def __init__(self, in_channels, out_channels, weight_init='normal', bias_init='zeros'):
        super(Dense, self).__init__()
        self.scalar = 1
        self.weight = Parameter(initializer(weight_init, [out_channels, in_channels]), name="weight")
        self.bias = Parameter(initializer(bias_init, [out_channels]), name="bias")
        self.matmul = ops.MatMul(transpose_b=True)

    def construct(self, x):
        x = self.matmul(x, self.weight)
        output = ops.bias_add(x, self.bias)
        return output

input_a = Tensor([[1, 2, 3], [4, 5, 6]], mindspore.float32)
input_b = Tensor([[2, 3, 4], [5, 6, 7]], mindspore.float32)
input_c = Tensor([[3, 4, 5], [6, 7, 8]], mindspore.float32)

dense_net = Dense(3, 4)
print(dense_net(input_a))
print(dense_net(input_b))
print(dense_net(input_c))

inputs = mnp.stack([input_a, input_b, input_c])

vmap_dense_net = vmap(dense_net)
print(vmap_dense_net(inputs))

Cell和函数式的自动向量化用法基本一致，只需要将vmap的第一个入参替换为Cell实例即可，Vmap将construct转换为作用于批处理数据的向量化construct。另外，该用例中初始化函数定义了两个Parameter参数, Vmap对于这类执行函数的自由变量的处理等同于将其作为入参并配置对应in_axes位置为None的场景。

通过这种方式，我们可以实现批量输入在同一个模型上进行训练或推理等功能，与现有网络模型输入支持batch轴输入的区别在于，利用Vmap实现的批处理维度更加灵活，不局限于NCHW等输入格式。

模型集成场景

模型集成场景将来自多个模型的预测结果组合在一起，传统的实现方式是通过分别在某些输入上运行各个模型，然后将各自的预测结果组合在一起。假如您正在运行的是具有相同架构的模型，那么您可以借助Vmap将它们进行向量化，从而实现加速效果。

该场景下涉及权重数据的向量化，如果您运行的模型是通过函数式编程形式实现，即权重参数在模型外部定义并通过入参传递给模型操作，那您可以直接通过配置in_axes的方式进行相应的批处理。而MindSpore框架为了提供便捷的模型定义功能，绝大部分nn接口的权重参数都在接口内部定义并初始化，这意味着模型中的权重参数在原始Vmap中无法对权重进行批处理，改造成通过入参传递的函数式实现需要额外工作量。不过您不必担心，MindSpore的vmap接口已经替您优化了该场景。您只需要将运行的多个模型实例以CellList的形式传入给vmap，框架即可自动实现权重参数的批处理。

让我们演示如何使用一组简单的CNN模型来实现模型集成推理和训练。

class LeNet5(nn.Cell):
    """
    LeNet-5网络结构
    """
    def __init__(self, num_class=10, num_channel=1):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(num_channel, 6, 5, pad_mode='valid')
        self.conv2 = nn.Conv2d(6, 16, 5, pad_mode='valid')
        self.fc1 = nn.Dense(16 * 5 * 5, 120)
        self.fc2 = nn.Dense(120, 84)
        self.fc3 = nn.Dense(84, num_class)
        self.relu = nn.ReLU()
        self.max_pool2d = nn.MaxPool2d(kernel_size=2, stride=2)
        self.flatten = nn.Flatten()

    def construct(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.conv2(x)
        x = self.relu(x)
        x = self.max_pool2d(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.relu(x)
        x = self.fc3(x)
        return x

假设我们正在验证同一模型架构在不同权重参数下的效果，让我们模拟四个已经训练好的模型实例和一份batch大小为16，尺寸为32 x 32的虚拟图像数据集的minibatch。

net1 = LeNet5()
net2 = LeNet5()
net3 = LeNet5()
net4 = LeNet5()

minibatch = Tensor(mnp.randn(3, 1, 32, 32), mindspore.float32)

相较于利用for循环分别运行各个模型后将预测结果集合到一起，Vmap能够一次运行获得多个模型的预测结果。

总结

本教程重点在于介绍Vmap的场景使用说明，本质上自动向量化并非将循环逻辑执行于函数外部，而是将循环下沉至函数的各个基元操作中，并将映射轴信息在基元操作间传递，从而保证计算逻辑的正确性。Vmap的性能收益主要来自于各个基元操作所对应的VmapRule实现，由于循环下沉至算子层级，因而更容易结合并行技术进行性能优化，如果您有自定义算子的场景也可以尝试为自定义算子实现特定的VmapRule，从而获得更好的性能。对于性能极致追求的场景还可以再结合图算融合特性进行优化。