搜索
首页科技周边人工智能谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

谷歌力推的JAX在最近的基准测试中性能已经超过Pytorch和TensorFlow,7项指标排名第一。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

而且测试并不是在JAX性能表现最好的TPU上完成的。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

虽然现在在开发者中,Pytorch依然比Tensorflow更受欢迎。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

但未来,也许有更多的大模型会基于JAX平台进行训练和运行。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

模型

最近,Keras团队为三个后端(TensorFlow、JAX、PyTorch)与原生PyTorch实现以及搭配TensorFlow的Keras 2进行了基准测试。

首先,他们为生成式和非生成式人工智能任务选择了一组主流的计算机视觉和自然语言处理模型:

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

对于模型的Keras版本,其采用了KerasCV和KerasNLP中已有的实现进行构建。而对于原生的PyTorch版本,则选择了网络上最流行的几个选项:

- 来自HuggingFace Transformers的BERT、Gemma、Mistral

- 来自HuggingFace Diffusers的StableDiffusion

- 来自Meta的SegmentAnything

他们将这组模型称作「Native PyTorch」,以便与使用PyTorch后端的Keras 3版本进行区分。

他们对所有基准测试都使用了合成数据,并在所有LLM训练和推理中使用了bfloat16精度,同时在所有LLM训练中使用了LoRA(微调)。

根据PyTorch团队的建议,他们在原生PyTorch实现中使用了torch.compile(model, mode="reduce-overhead")(由于不兼容,Gemma和Mistral训练除外)。

为了衡量开箱即用的性能,他们使用高级API(例如HuggingFace的Trainer()、标准PyTorch训练循环和Keras model.fit()),并尽可能减少配置。

硬件配置

所有基准测试均使用Google Cloud Compute Engine进行,配置为:一块拥有40GB显存的NVIDIA A100 GPU、12个虚拟CPU和85GB的主机内存。

基准测试结果

表2显示了基准测试结果(以步/毫秒为单位)。每步都涉及对单个数据批次进行训练或预测。

结果是100步的平均值,但排除了第一个步,因为第一步包括了模型创建和编译,这会额外花费时间。

为了确保比较的公平性,对于相同的模型和任务(不论是训练还是推理)都使用相同的批大小。

然而,对于不同的模型和任务,由于它们的规模和架构有所不同,可根据需要调整数据批大小,从而避免因过大而导致内存溢出,或是批过小而导致GPU使用不足。

过小的批大小也会使PyTorch看起来较慢,因为会增加Python的开销。

对于大型语言模型(Gemma和Mistral),测试时也使用了相同的批处理大小,因为它们是相同类型的模型,具有类似数量的参数(7B)。

考虑到用户对单批文本生成的需求,也对批大小为1的文本生成情况进行了基准测试。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

关键发现

发现1

不存在「最优」后端。

Keras的三种后端各展所长,重要的是,就性能而言,并没有哪一个后端能够始终胜出。

选择哪个后端最快,往往取决于模型的架构。

这一点突出了选择不同框架以追求最佳性能的重要性。Keras 3可以帮助轻松切换后端,以便为模型找到最合适的选择。

发现2

Keras 3的性能普遍超过PyTorch的标准实现。

相对于原生PyTorch,Keras 3在吞吐量(步/毫秒)上有明显的提升。

特别是,在10个测试任务中,有5个的速度提升超过了50%。其中,最高更是达到了290%。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

如果是100%,意味着Keras 3的速度是PyTorch的2倍;如果是0%,则表示两者性能相当

发现3

Keras 3提供一流的「开箱即用」性能。

也就是,所有参与测试的Keras模型都未进行过任何优化。相比之下,使用原生PyTorch实现时,通常需要用户自行进行更多性能优化。

除了上面分享的数据,测试中还注意到在HuggingFace Diffusers的StableDiffusion推理功能上,从版本0.25.0升级到0.3.0时,性能提升超过了100%。

同样,在HuggingFace Transformers中,Gemma从4.38.1版本升级至4.38.2版本也显著提高了性能。

这些性能的提升凸显了HuggingFace在性能优化方面的专注和努力。

对于一些手动优化较少的模型,如SegmentAnything,则使用了研究作者提供的实现。在这种情况下,与Keras相比,性能差距比大多数其他模型更大。

这表明,Keras能够提供卓越的开箱即用性能,用户无需深入了解所有优化技巧即可享受到快速的模型运行速度。

发现4

Keras 3的表现始终优于Keras 2。

例如,SegmentAnything的推理速度提升了惊人的380%,StableDiffusion的训练处理速度提升了150%以上,BERT的训练处理速度也提升了100%以上。

这主要是因为Keras 2在某些情况下直接使用了更多的TensorFlow融合操作,而这可能对于XLA的编译并不是最佳选择。

值得注意的是,即使仅升级到Keras 3并继续使用TensorFlow后端,也能显著提升性能。

谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择

结论

框架的性能在很大程度上取决于具体使用的模型。

Keras 3能够帮助为任务选择最快的框架,这种选择几乎总能超越Keras 2和PyTorch实现。

更为重要的是,Keras 3模型无需进行复杂的底层优化,即可提供卓越的开箱即用性能。

以上是谷歌狂喜:JAX性能超越Pytorch、TensorFlow!或成GPU推理训练最快选择的详细内容。更多信息请关注PHP中文网其他相关文章!

声明
本文转载于:51CTO.COM。如有侵权,请联系admin@php.cn删除
及时工程中的思想图是什么及时工程中的思想图是什么Apr 13, 2025 am 11:53 AM

介绍 在迅速的工程中,“思想图”是指使用图理论来构建和指导AI的推理过程的新方法。与通常涉及线性S的传统方法不同

优化您的组织与Genai代理商的电子邮件营销优化您的组织与Genai代理商的电子邮件营销Apr 13, 2025 am 11:44 AM

介绍 恭喜!您经营一家成功的业务。通过您的网页,社交媒体活动,网络研讨会,会议,免费资源和其他来源,您每天收集5000个电子邮件ID。下一个明显的步骤是

Apache Pinot实时应用程序性能监视Apache Pinot实时应用程序性能监视Apr 13, 2025 am 11:40 AM

介绍 在当今快节奏的软件开发环境中,确保最佳应用程序性能至关重要。监视实时指标,例如响应时间,错误率和资源利用率可以帮助MAIN

Chatgpt击中了10亿用户? Openai首席执行官说:'短短几周内翻了一番Chatgpt击中了10亿用户? Openai首席执行官说:'短短几周内翻了一番Apr 13, 2025 am 11:23 AM

“您有几个用户?”他扮演。 阿尔特曼回答说:“我认为我们上次说的是每周5亿个活跃者,而且它正在迅速增长。” “你告诉我,就像在短短几周内翻了一番,”安德森继续说道。 “我说那个私人

pixtral -12b:Mistral AI'第一个多模型模型 - 分析Vidhyapixtral -12b:Mistral AI'第一个多模型模型 - 分析VidhyaApr 13, 2025 am 11:20 AM

介绍 Mistral发布了其第一个多模式模型,即Pixtral-12b-2409。该模型建立在Mistral的120亿参数Nemo 12B之上。是什么设置了该模型?现在可以拍摄图像和Tex

生成AI应用的代理框架 - 分析Vidhya生成AI应用的代理框架 - 分析VidhyaApr 13, 2025 am 11:13 AM

想象一下,拥有一个由AI驱动的助手,不仅可以响应您的查询,还可以自主收集信息,执行任务甚至处理多种类型的数据(TEXT,图像和代码)。听起来有未来派?在这个a

生成AI在金融部门的应用生成AI在金融部门的应用Apr 13, 2025 am 11:12 AM

介绍 金融业是任何国家发展的基石,因为它通过促进有效的交易和信贷可用性来推动经济增长。交易的便利和信贷

在线学习和被动攻击算法指南在线学习和被动攻击算法指南Apr 13, 2025 am 11:09 AM

介绍 数据是从社交媒体,金融交易和电子商务平台等来源的前所未有的速度生成的。处理这种连续的信息流是一个挑战,但它提供了

See all articles

热AI工具

Undresser.AI Undress

Undresser.AI Undress

人工智能驱动的应用程序,用于创建逼真的裸体照片

AI Clothes Remover

AI Clothes Remover

用于从照片中去除衣服的在线人工智能工具。

Undress AI Tool

Undress AI Tool

免费脱衣服图片

Clothoff.io

Clothoff.io

AI脱衣机

AI Hentai Generator

AI Hentai Generator

免费生成ai无尽的。

热门文章

R.E.P.O.能量晶体解释及其做什么(黄色晶体)
3 周前By尊渡假赌尊渡假赌尊渡假赌
R.E.P.O.最佳图形设置
3 周前By尊渡假赌尊渡假赌尊渡假赌
R.E.P.O.如果您听不到任何人,如何修复音频
3 周前By尊渡假赌尊渡假赌尊渡假赌
WWE 2K25:如何解锁Myrise中的所有内容
4 周前By尊渡假赌尊渡假赌尊渡假赌

热工具

mPDF

mPDF

mPDF是一个PHP库,可以从UTF-8编码的HTML生成PDF文件。原作者Ian Back编写mPDF以从他的网站上“即时”输出PDF文件,并处理不同的语言。与原始脚本如HTML2FPDF相比,它的速度较慢,并且在使用Unicode字体时生成的文件较大,但支持CSS样式等,并进行了大量增强。支持几乎所有语言,包括RTL(阿拉伯语和希伯来语)和CJK(中日韩)。支持嵌套的块级元素(如P、DIV),

WebStorm Mac版

WebStorm Mac版

好用的JavaScript开发工具

VSCode Windows 64位 下载

VSCode Windows 64位 下载

微软推出的免费、功能强大的一款IDE编辑器

EditPlus 中文破解版

EditPlus 中文破解版

体积小,语法高亮,不支持代码提示功能

螳螂BT

螳螂BT

Mantis是一个易于部署的基于Web的缺陷跟踪工具,用于帮助产品缺陷跟踪。它需要PHP、MySQL和一个Web服务器。请查看我们的演示和托管服务。