为 JAX 用户提供模型检查点保存与恢复的通用工具库。
541+1 本周109PythonApache-2.0更新于 22小时前
google-research/kauldron
基于 JAX 和 Flax 的模块化机器学习模型训练库，面向快速科研迭代。
- Star
- 358
- Fork
- 36
- 未关闭 Issue
- 83
- 最近提交
- 昨天
这是什么？
Kauldron 是一个用于训练机器学习模型的 Python 库，设计重点是科研迭代速度和模块化。其各组件相互独立，可以脱离框架单独使用，便于与自有数据管线、Flax 模型和优化器配合。
它将 TFDS、Grain、SeqIO 等数据集，任意 Flax 模型和优化器衔接在一起。配置可以自由定制和覆盖，例如扫描不同模型结构，或修改内部某一层的参数。仓库提供了 MNIST 自编码器等示例配置，可作为起点。
该库支持以 Colab 为主的工作流，便于快速原型开发。内置功能包括绘图集成、性能分析器、事后调试以及运行时形状检查。
核心功能
- 01兼容 TFDS、Grain、SeqIO 或自定义数据管线。
- 02可接入任意 Flax 模型和优化器。
- 03配置可在任意层级进行扫描和覆盖。
- 04组件可脱离框架单独使用。
- 05内置性能分析器和运行时形状检查。
适用场景
- 在 Colab 笔记本中快速原型验证新模型结构。
- 为机器学习研究运行超参数或模型结构扫描。
- 使用统一的训练配置在自定义数据集上训练 Flax 模型。
使用前须知
- README 声明这不是 Google 官方支持的产品。
- README 没有提供安装或运行说明。
标签
常见问题
- kauldron 是什么？
- 基于 JAX 和 Flax 的模块化机器学习模型训练库，面向快速科研迭代。 Kauldron 是一个用于训练机器学习模型的 Python 库，设计重点是科研迭代速度和模块化。其各组件相互独立，可以脱离框架单独使用，便于与自有数据管线、Flax 模型和优化器配合。
- kauldron 可以免费使用吗？用的什么许可证？
- kauldron 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。
- kauldron 还在维护吗？
- 活跃维护中。kauldron 最近一次提交是在 2026年11月6日。
- kauldron 适合用来做什么？
- 在 Colab 笔记本中快速原型验证新模型结构。 为机器学习研究运行超参数或模型结构扫描。 使用统一的训练配置在自定义数据集上训练 Flax 模型。
- 在哪里获取 kauldron？
- 源代码和安装说明在 GitHub 上：https://github.com/google-research/kauldron
概要由 AI 根据项目 README 生成；GitHub 数据每日更新。
相关项目
基于 JAX/Flax 的 Vision Transformer 与 MLP-Mixer 图像分类模型代码及预训练权重。
1.3万+8 本周1478Jupyter NotebookApache-2.0更新于 上个月
Keras 3 是一个多后端深度学习框架，可运行于 JAX、TensorFlow、PyTorch 之上，OpenVINO 仅用于推理。
6.4万+21 本周2万PythonApache-2.0维护活跃更新于 前天
一个 Python 机器学习代码转换库，可在 PyTorch、TensorFlow、JAX、NumPy 等框架之间互相转换代码。
1.4万5437PythonOther更新于 17天前
面向 JAX 的神经网络库，强调灵活性，提供 Flax NNX API 用于构建、检查和调试模型。
7366+4 本周849Jupyter NotebookApache-2.0更新于 前天
一键安装的 JAX 工具包，打包了经过集成测试、版本兼容的 Flax、Optax、Orbax 等 AI 与数值计算库。
31551PythonApache-2.0更新于 25天前