AXLearn 深度学习库
此库正在积极开发中,API 可能会有所变化。
目录
介绍
AXLearn 是建立在 JAX 和 XLA 之上的库,支持大规模深度学习模型的开发。
AXLearn 采用面向对象的方法来应对从构建、迭代到维护模型过程中遇到的软件工程挑战。 该库的配置系统允许用户从可重用的构建块中组合模型,并与其他库(如 Flax 和 Hugging Face transformers)集成。
AXLearn 为扩展而生。 它支持在数以千计的加速器上高效利用训练多达数千亿参数的模型。 它还设计为能够在公共云上运行,并提供工具来部署和管理作业和数据。 基于 GSPMD 构建,AXLearn 采用全球计算范式,让用户能够在虚拟全球计算机上而不是单个加速器上描述计算。
AXLearn 支持广泛的应用,包括自然语言处理、计算机视觉和语音识别,并且包含训练最先进模型的基线配置。