用于从 TensorFlow 到 JAX 的深度学习模型迁移的多智能体 AI 系统

人工智能产品及其底层模型的快速发展带动了深度学习框架的不断创新。谷歌一直在数十种产品中率先使用机器学习。在不同的机器学习框架和版本中维护大量模型源代码是一项重大挑战。到目前为止，维护和迁移工作主要由人类专家手动完成。我们描述了一个基于 AI 的多代理系统，该系统是为了支持基于 TensorFlow 的深度学习模型自动迁移到基于 JAX 的模型而构建的。我们做出了三个主要贡献：首先，我们展示了结合使用静态分析和 AI 指令的 AI 规划器如何使用 AI 生成的基于示例的剧本，为非常复杂的代码组件创建迁移计划，并由编排器和编码器的组合可靠地遵循这些计划。其次，我们定义质量指标和基于人工智能的判断器，当要评估的代码没有测试并且必须遵守严格的样式和依赖性要求时，它们可以加速开发。第三，我们演示了系统如何在商业现实用例的大型超大规模环境中加速代码迁移。我们的方法极大地减少了深度学习模型迁移的时间（加速 6.4 倍至 8 倍），并创建了一个良性循环，让 AI 有效支持其自身的开发工作流程。我们期望这里描述的技术和方法可以推广到其他框架迁移和一般代码转换任务。