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消耗内存大,而且没有原地操作,所以在小模型上浪费时间。几年前在Netflix的时候,他设计开发了一个D语言框架vectorflow。目前,他在github上获得了1200颗星。机器学习模型的假设SimpleChains.jl是Pumas-AI和JuliaComputing、Roche和马里兰大学巴尔的摩分校合作开发的一个库,其主要目标是为小型神经网络提供尽可能高的性能。SimpleChains.jl最初是作为医学数据分析中科学机器学习(SciML)的解决方案:小型神经网络(以及其他逼近器,例如傅里叶级数或切比雪夫多项式展开)可以与已知的半生理模型相结合,以前未知的机制和预后因素被发现。从黑洞动力学到抗震安全建筑的开发,SciML方法的有效性已在许多学科中得到证明,可以灵活地发现/指导(生物)物理方程。应用场景变化太大。在这种情况下,可以通过使用一些专门化的神经网络来提高模型的性能。具体来说,在机器学习模型的研究中,通常依赖于一个假设:神经网络足够大,矩阵乘法(如卷积)的O(n^3)时间成本占了大部分运行时间,这意味着基本上,它也是机器学习库的大多数机制背后的4个指导原则:1.矩阵乘法的复杂度是三次的,内存分配的大小是线性的,所以使用非分配内存。操作向量的优先级不高;2、目前AI加速的工作主要集中在GPU核心加速,让指令尽可能快的跑起来,因为这些大型的矩阵-矩阵运算在GPU上是最快的,而且也用于大型模型.主要瓶颈,所以性能基准基本上只是测量这些特定内核的速度;3、做自动差分反向传播时,几乎看不到将值复制到内存的操作,内存分配被更大的内核调用隐藏起来;4..用户可以自由编写磁带生成反向传播。虽然它增加了前向过程中构建字典的成本,但它也会被更大的内核调用所覆盖。然而,这些假设在实际案例中真的成立吗?如果没有,是否可以着重改进这些方面来提高计算性能?小型神经网络的瓶颈在哪里?对于初学者,可以先测试假设1和假设2,用一段Julia代码测试内存申请时间、GPU计算时间等。可以看出,当我们进行大型矩阵乘法运算时,比如100x100*100x100,基本上可以忽略由于内存分配引起的任何开销。但也可以看出,在低端可以看到一些相当显着的性能提升,这些提升是通过使用纯JuliaLoopVectorization.jl实现的,因为标准BLAS工具往往在这个区域有额外的线程开销(再次,这方面没有优化)。如果您一直在利用GPU带来的好处而没有深入研究细节,那么这个事实可能会让您大吃一惊!GPU被设计成具有许多内核的慢速芯片,因此它们仅对非常并行的操作有效,例如大型矩阵乘法。正是从这一点来看,假设2可以被认为是一个大型网络操作。但同样,在小型网络的情况下,由于缺乏并行计算,使用GPU内核的性能可能不如设计良好的CPU内核。矩阵运算只有在可以使用批处理时才会发生(A*B中B矩阵的每一列都是一个单独的批处理)。在大多数科学机器学习环境中,例如计算ODE邻接中的向量雅可比积,此操作是矩阵向量乘法。这些操作的时间复杂度仅为O(n^2),这样的话内存开销会被放大。神经网络的基本运算是Sigma,所以还有一个时间复杂度为O(n)的运算。这种情况下,内存开销就更严重了。对于假设3和4,需要更多关注反向传播的实现。不同机器学习库的自动微分方法也存在差异。有些库立即反向传播梯度值,有些库需要保存梯度,这需要额外的内存开销操作。在具体应用中,如果已知梯度立即传播,则可以立即计算梯度。与一般的实现相比,只需要一个缓存向量就可以解决问题,并且赋值到位,自动微分的额外开销都没有了。基于这些想法,研究人员开源了SimpleChains.jl,可以很好的解决这类优化问题,并且可以在CPU上快速拟合和优化小模型。大多数早期的神经网络原型模型设计希望:1.实现更好2.专注于小尺寸模型,在早期开发阶段放弃一些针对大模型的内核优化操作(如缓存平铺);3.有一个API,其中向量的参数和梯度都是一流的,这样更容易与各种优化器或求解器(如BFGS)一起工作;4、使用“纯Julia”编写,更方便开发和优化;在大量使用LoopVectorization.jl的同时,SimpleChains.jl不依赖于任何BLAS或NN库。开发人员的长期目标是扩展此循环编译器优化以自动生成回调。但是这种以编译器为中心的方法已被用于实现方便:虽然我们仍然需要手写梯度,但我们不需要手动优化它们。SimpleChains.jl的实际性能如何?研究人员用2×2矩阵做了一个实验,并在英特尔i9-10980XE上用AVX512指令集运行。10,000个epoch只用了0.41秒,而pyTorch只用了15秒,据说在这个微型神经网络上,加速大约是35倍。将实验切换到带有AVX2指令的AMDEPYC7513机器上,Julia的实现耗时0.72秒,而PyTorch的实现耗时70秒,差距扩大到100倍。研究人员还在AMDRyzen95950X上试验了Jax代码。Julia用了1.3秒,Jax用了14秒,增长了大约10倍。切换到Intel(R)Core(TM)i9-10980XECPU@3.00GHz平台,Jax需要9秒,Julia需要0.4秒,大约快22倍。然后换个更差的处理器,在6核CPU上,Jax需要19秒,而Julia需要9秒,速度提升只有2倍。在稍微大一点的、实际可用的神经网络上,训练速度还会有这么大的差距吗?研究人员使用LeNet5来测试MNIST。这个例子只是一个非常保守的速度估计,因为在更传统的机器学习用例中,批处理可以使用矩阵乘法,但即使在这种情况下,由于网络规模半小,也能看到可观的性能优势。训练10个epoch,batchsize为2048,使用PyTorch在A100上训练两次耗时17.66和17.62,准确率分别为94.91%和96.92%;在V100上的训练时间分别为16.29和15.94,准确率分别为95.6%和97.5%。然而,这个问题对于GPU来说仍然是大材小用。在2048的batchsize上计算速度还是很快的,时间主要花在了CPU转移到GPU上。在AMDEPYC7513和Inteli910980XE上又进行了两次实验,事实证明它们比GPU更快、更准确。切换到SimpleChains.jl,AMD平台耗时3秒,准确率98.3%;在Intel平台上,仅需1秒,准确率达98.2%;即使在笔记本的Intel平台上,耗时也只有5.3秒,准确率为97%目前的大型机器学习框架做的非常好,专注于为99.9%的用户提供一流的性能,但在手中另外0.1%的小模型用户,框架不好用。这就是可组合性和灵活性的美妙之处:一种允许您轻松构建机器学习框架的语言也是一种允许您构建针对替代人群优化的替代框架的语言。
