Julia开源新框架SimpleChain：小型神经网络速度比PyTorch快5倍！

---

  新智元报道  

编辑：LRS

【新智元导读】世上没有免费的午餐，享受了通用框架的便利，在特定任务上就要牺牲性能。最近Julia开源了一个新框架SimpleChain，在小型神经网络的运行速度上比PyTorch至少快5倍！

Julia从一出生开始，就瞄准了科学计算领域，并且一直在与Python暗中较量。

 

在神经网络的框架上，Python有PyTorch和TensorFlow，几乎是深度学习开发的首选框架，并且获得了Meta和Google在技术和资金上的支持，蓬勃发展。

 

虽然Julia也有Flux.jl框架，但Julia社区一直依赖于语言本身的高性能产生的生产力，所以Flux.jl的代码量相比Python框架来说，可以称得上是特别「苗条」了，例如PyTorch和TensorFlow包括了整个独立的语言和编译器（torchscript、XLA等），而Flux.jl仅仅由Julia语言编写。

 

 

当然，世界上没有免费的午餐，如果以不同的视角来看，想要在机器学习领域开发出一个简单、通用且高性能的框架几乎是不可能的，只能不断权衡。

 

比如对于一个特定的问题，如果需要稀疏的小模型，想要获得最高性能的方法就是重写一遍，而非采用通用框架。

 

最近Julia社区又开源了一个新框架SimpleChains.jl，在小模型场景下相比PyTorch最少能提速5倍。

 

代码链接：https://github.com/PumasAI/SimpleChains.jl

 

开发人员表示，这个框架不会对所有人都有用，但对那些需要它的人来说，它是非常有用的。

 

有网友表示十分赞同：「不同的任务用不同的工具」，因为TF和pyTorch消耗了大量的内存，并且没有原地操作，所以在小模型上很浪费时间。他数年前在Netflex时就设计开发了一个D语言的框架vectorflow，目前在github已获1200个stars

 

 

机器学习模型的假设

SimpleChains.jl是由Pumas-AI和Julia Computing与Roche和马里兰大学巴尔的摩分校合作开发的一个库，它的主要目的就是为小型神经网络提供尽可能高的性能。

 

SimpleChains.jl最开始用于在医疗数据分析中用于科学机器学习（SciML）的解决方案：小型神经网络（和其他近似器，如傅里叶数列或切比雪夫多项式展开）可以与已知的半生理学模型（semi-physiologic models）相结合，发现以前未知的机制和预后因素。

 

从黑洞动力学到地震安全建筑的开发，SciML方法的有效性已经在许多学科中得到证实，能够灵活地发现/指导（生物）物理方程。

 

应用场景变化太大，在这种情况下，使用一些专用（specialization）的神经网络才有可能提升模型的运行性能。

 

 

具体来说，在机器学习模型的研究中，通常依赖于一个假设：神经网络足够大，其中矩阵乘法（如卷积）的O(n^3)时间成本占了运行时间的绝大部分，这基本上也是机器学习库的大部分机制背后的4大指导原则：

 

1. 矩阵乘法的复杂度是立方的，而内存分配的规模是线性的，所以用非分配（non-allocating）内存的方式来操作向量的优先级并不高；

 

2. 目前AI加速的工作主要集中于GPU内核加速，让指令运行尽可能快，由于这些大型矩阵-矩阵操作在GPU上是最快的，并且也是大模型的主要瓶颈，所以性能基准基本上只是衡量这些特定内核的速度；

 

3. 当做自动微分反向传播时，将数值复制到内存的操作几乎感觉不到，内存分配被较大的内核调用所隐藏；

 

4. 用户可以随意写一个tape来生成反向传播，虽然增加了在前向过程中建立字典的成本，但是也会被更大的内核调用所掩盖。

 

但，这些假设在真实的案例中是否真的能全部成立？

 

如果不成立的话，能不能把重点放在这些方面的改进，从而提升更高的运算性能？

 

小型神经网络的瓶颈在哪？

对于初学者来说，可以先测试一下假设1和2，通过一段Julia代码来测试内存申请时间、GPU运算时间等。

 

 

可以看到当我们进行较大的矩阵乘法操作时，比如100x100*100x100，基本可以忽略由于内存分配而产生的任何开销。

 

 

但同时也可以看到，在lower end有可能出现一些相当显著的性能提升，这些收益是通过使用纯Julia LoopVectorization.jl实现的，因为标准的BLAS工具在这个区域往往有额外的线程开销（同样，在这个区域没有进行优化）。

 

如果你一直在利用GPU带来的好处而不去研究细节，那么这个事实可能会让你大吃一惊！GPU被设计成具有许多内核的慢速芯片，因此它们只对非常并行的操作有效，例如大型矩阵乘法。正是从这一点出发，假设2可以被认为是大型网络操作。

 

但同样，在小网络的情况下，由于缺乏并行计算，使用GPU内核的性能可能还不如设计良好的CPU内核。

 

矩阵操作只有在能够使用批处理（A*B中的B矩阵的每一列都是一个单独的批处理）时才会发生。

 

在大部分科学机器学习的情境下，如ODE邻接中的向量Jacobian乘积的计算，这种操作是矩阵-向量乘法。这些操作的时间复杂度只有O(n^2)，在这种情况下内存开销会被放大。

 

 

神经网络的基本操作是Sigma，所以还有一个O(n)时间复杂度的操作，这种情况下内存开销显得更严重。

 

 

对于假设3和4来说，需要更加关注反向传播的实现。不同机器学习库的自动微分方法也存在着区别。有些库是立刻反向传播梯度值，也有些需要把梯度保存起来，这样就又需要额外的内存开销操作了。

 

 

在特定的应用里面，如果知道梯度立刻传播，就可以立即计算梯度，相比通用实现来说，只需要一个缓存向量解千愁，原地赋值，这样的话所有自动微分的额外开销都没有了。

 

 

基于这些想法，研究人员开源了SimpleChains.jl，可以很好地解决这类优化问题，可以在CPU上快速拟合和优化小模型，早期的神经网络原型模型设计大多都希望：

 

1. 达到更好的性能，最好能达到CPU的峰值FLOPs；

2. 专注于小尺寸的模型，在早期开发阶段放弃一些针对大型模型的内核优化操作（如缓存平铺）；

3. 有一个API，其中的向量的参数和梯度都是first class，以便更容易地与各种优化器或求解器（如BFGS）协同工作；

4. 使用「纯Julia」编写，更方便开发和优化；在大量使用LoopVectorization.jl的同时，SimpleChains.jl并不依赖任何BLAS或NN库。

 

开发人员的长期目标是将这种循环编译器的优化方法扩展到自动产生pullbacks。但这种以编译器为中心的方法已经被用于实现的便利性：虽然我们仍然需要手写梯度，但我们不需要对它们进行手工优化。

 

SimpleChains.jl的实际性能怎么样？

 

研究人员用一个2×2的矩阵做了一个实验，在带有AVX512指令集的Intel i9-10980XE跑了一下，1万个epoch花了0.41秒，相比之下pyTorch花了15秒，也就是说在这种微型神经网络上，提速大约35倍。

 

 

把实验换到AMD EPYC 7513 带有AVX2指令的机器上，Julia的实现花费时间为0.72秒，而PyTorch的实现则需要70秒，差距拉升到了100倍。

 

研究人员又在AMD Ryzen 9 5950X实验了一份Jax代码，Julia耗时为1.3秒，Jax则需要14秒，提升约10倍。

 

换到Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz 平台上，Jax耗时为9秒，Julia需要0.4秒，大约22倍提升。

 

再换到差一点的处理器，6核CPU上，Jax需要19秒，而Julia需要9秒，速度提升就只有2倍了。

 

在稍微大一点的、实际可用的神经网络上，训练速度还会有这么大的差距吗？

 

研究人员用LeNet5来测试MNIST，这个例子只是一个非常保守的速度估计，因为在更传统的机器学习用例中，批处理可以使用矩阵乘法，不过即使在这种情况下，由于semi-small的网络规模，也能看到大量的性能优势。

 

 

在batch size为2048的情况下训练10个epoch，用PyTorch在A100上训练两次耗时为17.66和17.62，准确率分别为94.91%和96.92%；在V100上训练时间为16.29和15.94，准确率分别为95.6%和97.5%

 

不过这个问题对于GPU来说还是杀鸡用牛刀了，在2048的batch size上运算速度还是很快，时间主要耗费在CPU转移到GPU上了。

 

在AMD EPYC 7513和Intel i9 10980XE又进行了两次实验，结果比GPU更快，准确率也更高。

 

 

换到SimpleChains.jl，在AMD平台上耗时为3秒，准确率98.3%；在Intel平台上，耗时仅为1秒，准确率为98.2%；即使在笔记本的Intel平台上，耗时也仅为5.3秒，准确率97%

 

目前大型机器学习框架在专注于为其99.9%的用户提供一流的性能方面做得非常好，但在另外0.1%的小模型用户手里，框架却不好用。

 

这就是可组合性和灵活性的优势：一种允许你轻松构建机器学习框架的语言，也是一种允许你构建替代框架的语言，这些框架针对替代人群进行优化。

参考资料：

https://julialang.org/blog/2022/04/simple-chains/