大家好,我是刘明,明志科技创始人,华为昇思MindSpore布道师。
技术上主攻前端开发、鸿蒙开发和AI算法研究。
努力为大家带来持续的技术分享,如果你也喜欢我的文章,就点个关注吧
正文开始
AI融合计算的蓬勃发展,对框架能力提出了新的需求和挑战。问题场景和模型设计逐渐复杂化,使得业务数据的维度和计算逻辑的嵌套深度也相应增长。结合向量化优化手段可以有效优化性能瓶颈,但实现向量化优化对于普通用户并非易事。虽然用户可以很容易地实现低维数据运算逻辑,但随着数据维度的增长,业务逻辑也变得更为复杂,用户需要清晰了解各操作间的数据维度的逻辑映射关系,给用户的模型设计和编码带来了巨大挑战。自动向量化特性(Vmap)帮助用户解决了这个头疼的问题,该技术允许用户将特定的批处理逻辑从函数中剥离。用户在编写函数时,只需要先考虑低维的运算逻辑即可,通过调用vmap接口自动实现高维运算,并且支持嵌套调用,有效降低问题复杂度。
本教程介绍自动向量化vmap接口的使用方式,将模型或函数中高度重复的运算逻辑转换为并行的向量运算逻辑,从而获得更加精简的代码逻辑以及更高效的执行性能。
向量化思维
向量化思维在提升计算性能的技术中十分常见,向量化思维可公式化表示为:
a
1
+
b
1
=
c
1
a
2
+
b
2
=
c
2
a
3
+
b
3
=
c
3
a
4
+
b
4
=
c
4
⇒
a
⃗
+
b
⃗
=
c
⃗
\begin{split}\begin{matrix} 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{matrix} \Rightarrow \vec{a} + \vec{b} = \vec{c}\end{split}
a1+b1=c1a2+b2=c2a3+b3=c3a4+b4=c4⇒a+b=c
其核心思想在于将多次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为我们提供了另一种实现方式。