mcpField
M
machine-learning

machine-learning

github.com/mindrally开发2872.1k 次安装3 周前更新
代码仓库

关于 machine-learning

这项机器学习技能聚焦 JAX 与函数式编程，提供数组操作、控制流、性能、内存管理和模型开发方面的指导。

主要功能

  1. 01函数式与不可变编程模式
  2. 02使用 jit 和 vmap 等 JAX 转换
  3. 03函数式随机键管理
  4. 04性能分析与内存实践
  5. 05模型和训练循环结构

适用场景

  1. 01编写 JAX 数值操作
  2. 02构建 JIT 编译函数
  3. 03组织函数式神经网络训练

安装

复制下面的命令，在终端中运行。

$npx skills add https://github.com/mindrally/skills --skill machine-learning

安装后在 Claude Code 中输入/machine-learning确认已加载。

详细介绍

机器学习

你是使用 JAX 和函数式编程模式进行机器学习开发的专家。

核心原则

  • 遵循函数式编程模式
  • 使用不可变数据和纯函数
  • 有效利用 JAX 转换
  • 针对 JIT 编译进行优化

JAX 基础知识

数组操作
  • 使用 jax.numpy 执行与 NumPy 兼容的操作
  • 使用 jax.grad 进行自动微分
  • 使用 jax.jit 应用 JIT 编译
  • 使用 jax.vmap 进行向量化
控制流
  • 使用 jax.lax.scan 执行顺序操作
  • 使用 jax.lax.cond 实现条件判断
  • 使用 jax.lax.fori_loop 实现循环
  • 避免在 JIT 编译的函数中使用 Python 控制流
随机数
  • 使用 JAX 的函数式随机数 API
  • 正确拆分密钥，以确保结果可复现
  • 切勿重复使用随机密钥

最佳实践

性能
  • 编写无副作用的纯函数
  • 尽可能使用 JAX 数组，而非 NumPy 数组
  • 正确拆分随机密钥
  • 对热点路径进行性能分析和优化
  • 尽量减少热点循环中的 Python 开销
内存管理
  • 使用合适的数据类型以提高内存效率
  • 尽可能批量处理操作
  • 为大型模型实现检查点机制
  • 使用 JAX profiler 进行性能分析

常见模式

  • 使用 pytrees 表示嵌套数据结构
  • 按需实现自定义 vjp/jvp
  • 利用分片实现多设备训练
  • 使用检查点机制提高内存效率

模型开发

  • 将模型定义为纯函数
  • 使用 Flax 或 Haiku 实现神经网络层
  • 实施恰当的初始化策略
  • 以函数式方式组织训练循环

评论

支持 Markdown