聊一聊AI框架前端

聊一聊AI框架的前端，涉及两块内容：

1、AI框架如何对接前端宿主编程语言

2、AI框架的编程范式

一、AI框架如何对接宿主编程语言

现在Python可以说是AI框架默认的宿主语言，开发者喜欢其易用性和灵活性，但是框架需要解决Python灵活性和性能的矛盾（如何进行Python加速）。

Python加速范式

在AI框架出来之前，广泛存在三种范式：

CPython：Python的C extension，现在主流的模式，完全开放Python解释器的内部数据接口和API，允许开发者使用Native语言编写扩展，直接访问这些数据接口；这样Python语言作为胶水语言提供灵活性，需要性能的地方通多Python的C extension进行加速。
Python JIT虚拟机：主要是期望在Python解释执行的基础上增加JIT编译加速的能力，典型的如PyPy；不过由于前期CPython暴露了太多内部接口，导致Python JIT虚拟机兼容的困难，也就是说CPython支撑了Python的成功，但是也阻止了Python JIT虚拟机的演进（来自我们语言虚拟机专家的观点），Pypy难说成功；
Python与JIT的混合模式：典型的如Numba，Python JIT虚拟机的一种妥协实现方式，通过修饰符，进行部分Python语句的加速。

AI框架也有类似的发展过程

早期：Pytorch是典型的CPython的范式，说透了很简单，但是Pytorch通过Tape模式解决了自动微分的问题，通过与GPU的有效协同（异步执行等）解决了性能的问题，同时又保持了Python解释器执行的灵活性（所谓的动态图），占得了先机；TF1.x的本质也是CPython，不过封装比较高层，改变了许多Python的原有使用习惯，变成了静态图执行方式，虽然性能提升更加明显，但是易用性和灵活性下降了很多。
后期：TorchScript、JAX、TF2.0，包括MindSpore自身，逐步走向CPython+Numba混合的模式，在动态图执行的时候，采用CPython模式，在静态图模式或者Staging模式下，通过模式设置或者修饰符方式，进行编译加速；不过AI框架的编译器和Numb编译有一定的区别，numba是从python直接lowering到机器码，AI的编译器是个分层的编译器，从图编译器——>算子编译器——>codegen逐步lowering，这样一方面既能借用CPython模式下的实现的算子能力，又能通过分层解决不同类型的挑战，减少系统复杂度。

未来的挑战和趋势

Python的编译加速很难获得完备性

目前AI框架进行Python编译加速主要两种方法：一种是Tracing；另外一种是AST转换。

Tracing的模式不好处理动态控制流；AST很难支持完备的Python语法。

本质的原因是Python这种解释器语言的动态语法对静态编译是不友好的

复杂性丢给了开发者

当前的模式需要开发者加修饰符进行性能加速，这就意味着开发者能够识别可以加速和需要加速的Python代码，门槛是比较高的。

未来解决之道的探讨：

TypePython：类似TypeScript一样，是否把当前的Python Type Hint做的更加易用一些/全面一点，如果这样的话，AI框架采用AOT或者JIT方式去执行Python至少可以做到比较完备。
自动JIT：不需要用户手工加修饰符去加速Python，系统自动进行JIT；LayzTensor的方式暂时还无法解决编译开销/缓存/Barrier时机等问题，PyPy这种标准的JIT方法是否更有效？
新的编程语言：也许前面说的问题都解决了，但是就怕没有开发者使用。

二、编程范式

AI框架的编程范式又很多分类的方法，比如动态图和静态图、命令式和声明式等，我下面想提另外一种分类：以函数为中心和以Tensor为中心。

以函数为中心：把神经网络看成一个复杂的函数；如JAX、MindSpore。
以Tensor为中心：把神经网络看成是一个dataflow的图；如Pytorch、TF2.0等。

看上去理念上差异比较大，但是对开发者来说，实际实现中，正向过程非常类似，因为函数也好/dataflow也好都是通过编程语言的函数调用来实现；而反向过程，双方就有一定的差异了。

对于函数式的风格来说，BP过程先对函数进行Gradient，得到BP函数（实际上得到的函数本身是一个正向和反向在一起的复合函数），然后在进行求值。

MindSpore为例：

#定义正向的网络/复合函数
class Net(nn.Cell):
     def __init__(self):
         super(Net, self).__init__()
         self.matmul = P.MatMul()
         self.z = Parameter(Tensor(np.array([1.0], np.float32)), name='z')
     def construct(self, x, y):
         x = x * self.z
         out = self.matmul(x, y)
         return out

#GradNet的功能是对输入的net进行Grad，返回一个带正反向的net
class GradNet(nn.Cell):
     def __init__(self, net):
         super(GradNet, self).__init__()
         self.net = net
         self.grad_op = GradOperation()
     def construct(self, x, y):
         gradient_function = self.grad_op(self.net)
         return gradient_function(x, y)

x = Tensor([[0.5, 0.6, 0.4], [1.2, 1.3, 1.1]], dtype=mstype.float32)
y = Tensor([[0.01, 0.3, 1.1], [0.1, 0.2, 1.3], [2.1, 1.2, 3.3]], dtype=mstype.float32)

#先得到BP函数，然后进行求值
NewNet = GradNet(Net())
output = NewNet(x,y)

#也可以把前面两步合并成一步
output = GradNet(Net())(x, y)

上面是带有深度模型的风格，还有纯函数的风格，以Jax为例：

grad_tanh = grad(jnp.tanh)
print(grad_tanh(2.0))

对于Tensor为中心的范式来说，BP过程其实拿到dataflow的执行结果的tensor，然后基于这个tensor进行反向传播，这个风格估计大家都很熟悉了，这里就不赘述。

以Pytorch为例：

#前向过程
y_pred = a + b * x + c * x ** 2 + d * x ** 3
#计算loss，loss是一个tensor
loss = (y_pred - y).pow(2).sum()
if t % 100 == 99:
    print(t, loss.item())

# 通过tensor进行bp
loss.backward()

这两种风格，各有优缺点：

函数式：符合算法的直观，除了深度学习场景外，也适合科学计算等，比如做高阶微分很方便，grad(grad())(....)；

Tensor的方式：非常符合深度学习的场景，直接使用tensor的结果，做计算过程的拼接相对方便，比函数式少一次封装。

未来编程范式的思考

现在AI框架的前端表达还没有完全收敛，框架各自的接口还是有差异的，但是总的来说，我比较喜欢JAX的风格，原因有两个：

Numpy+Scipy+Grad：我想Numpy和Scipy是事实上的标准，如果我们在计算和算法逻辑上能统一到这一块，也是比较自然。
分层解耦：JAX的基础包只提供简单的Numpy+Scipy+Grad的接口；一些高级库，如haiku，基于JAX基础包进行二次开发再提供高层的API，这就意味着大家如果在基础库上做到接口兼容就能进行高级库的使用。

原文：知乎
作者：金雪锋

推荐阅读

更多嵌入式AI技术干货请关注嵌入式AI专栏。

一、AI框架如何对接宿主编程语言

二、编程范式

推荐阅读

目录