Sunstreaker：基于 JAX 的轻量级深度学习框架

Sunstreaker：基于 JAX 的轻量级深度学习框架

sunstreaker 以jax为后端的类似keras的框架 项目地址: https://gitcode.com/gh_mirrors/su/sunstreaker

1. 项目基础介绍

Sunstreaker 是一个基于 JAX 的深度学习框架，旨在为研究人员和开发者提供一个简洁、高效的工具，以加速深度学习模型的开发与训练。该项目的编程语言主要是 Python。

2. 项目核心功能

• 类似 Keras 的 API：Sunstreaker 提供了类似 Keras 的接口，使得从 Keras 迁移到 Sunstreaker 变得非常容易。
• 基于 JAX 的后端：利用 JAX 的高效计算能力，Sunstreaker 可以在 CPU 和 GPU 上高效运行，并支持自动微分。
• 快速实验：项目设计用于快速实现新的想法和复现最新的研究成果。
• 分布式训练：支持快速分布式训练大型模型。
• 使用开源模型权重：方便地加载和使用开源模型的权重。

3. 项目最近更新的功能

• 内核改动：对 Layer 的 call 函数进行了优化，不再需要传入 params，build 输出也不再需要输出 params。
• Model params 变为有序字典：这样的改变方便了大模型参数的加载。
• build 不再需要输入随机种子：随机种子现在由内核自动分配，简化了开发流程。
• 新增 application layers：引入了 transformers/bert 相关的层，如 MultiHeadAttention、PositionEmbedding、FeedForward 等。
• 新增 application optimizers：支持了 diffusion/DDPM 和 AdamW 优化器。

Sunstreaker 的持续更新，为用户提供了更加强大和灵活的工具，以应对各种深度学习任务。

sunstreaker 以jax为后端的类似keras的框架 项目地址: https://gitcode.com/gh_mirrors/su/sunstreaker