Go机器学习框架之火重燃，Google前研究员开源期望媲美Jax的GoMLX

Go作为一门兼具高性能与简洁性的编程语言， 近年来在各种领域得到广泛应用 。然而，在 机器学习领域 ，Go相比Python、C++、Julia等语言，生态仍然较为薄弱。目前的Go机器学习框架无论在功能全面性上，还是在社区生态支持上都难以与TensorFlow、PyTorch、Jax等重量级框架抗衡。究其原因，笔者觉得还是 Go社区缺少熟悉和精通机器学习方面的人才 。

不过，随着时间的推移，总会有一些开发者会尝试探索在Go语言中实现高效机器学习的可能性。近期，笔者就发现一个由 Google前研究员Jan Pfeifer ^[1] 开源的 Go机器学习框架项目GoMLX ^[2] ，虽然我不是很了解机器学习这个领域，但也能看出一些该项目与以往Go机器学习项目的不同。

下面我就来简单介绍一下gomlx项目。

注：该项目开源时间不长，正处于初期开发阶段，还是一个实验性项目，也基本是由Jan Pfeifer一人完成。

一. GoMLX简介

GoMLX是一个主要由Go实现的、快速且易用的机器学习和通用数学库与工具集，你可以将其视为Go语言版本的PyTorch/Jax/TensorFlow(也许现在还不是，但其目标定位是这样的)。它力求成为Go语言中的标准化机器学习框架，让开发者在使用Go进行机器学习时，无需依赖其他语言或工具链。

结合该项目的README和作者提供的 GoMLX教程 ^[3] ，我们可以总结一下目前GoMLX目前已具备的主要功能特性，以及其设计上的优点。

GoMLX已具备的主要功能特性包括：

• 计算图构建：支持创建和操作计算图。
• JIT编译：能够即时编译计算图以提高执行速度。
• 多后端支持：可以使用不同的后端执行操作，默认使用XLA。
• 张量操作：提供了基本的张量操作和管理功能。
• 自动微分：支持静态梯度计算。
• 变量管理：通过Context对象管理变量和模型参数。
• 机器学习层：提供了多种常用的机器学习层，如Dense、Dropout、BatchNorm等。
• 训练工具：包括Trainer和Loop等用于模型训练的工具。
• 数据集接口：定义了Dataset接口用于数据输入。
• 优化器、损失函数和指标：提供了常用的优化器、损失函数和评估指标。
• 可视化支持：能够绘制训练过程中的指标图表。
• 调试工具：提供了多种调试方法，如形状断言、错误处理等。
• 预训练模型支持：提供了诸如InceptionV3等预训练模型。

GoMLX在设计上相对于其他Go机器学习框架的优点有以下几点：

• 模块化设计：将不同功能封装在独立的包中，便于理解和使用。
• 性能：通过使用JIT编译和XLA后端，实现了高性能计算。
• 错误处理：采用异常机制处理错误，提供了详细的堆栈跟踪，有助于调试。
• 文档完善：强调了文档的重要性，提供了详细的API文档和教程。
• 上下文管理：通过Context对象统一管理变量、参数和作用域，简化了模型开发。
• 硬件适配：支持CPU和GPU，并可能扩展到其他硬件（如TPU）。
• 开发体验：提供了交互式开发环境（如Jupyter notebook支持，通过 GoNB项目 ^[4] ，作者的另外一个开源项目），方便实验和学习。
• 可视化集成：内置了训练过程的可视化工具，方便监控和分析。

二. GoMLX的长期目标

作者为GoMLX项目设定了明确的长期目标，包括：

• 在Go中构建和训练模型，注重简单易读、API清晰、可组合性强。
• 成为高效的研究和教育平台，方便实验新的ML想法
• 成为可靠的生产平台，支持：
• 现代加速硬件如TPU和GPU
• 与TensorFlow Serving等工业工具集成
• 从Hugging Face Hub等导入预训练模型
• 将模型编译为二进制库或WebAssembly

三. GoMLX的TODO list

前面说过，该项目处于初期开发阶段，还有很多待开发的功能，作者在项目中放置了一个todo list，我们也可以看一下作者对GoMLX的功能规划：

1. 模型相关

• 增加CNN相关功能，如SeparableConv2D，以支持完整的ResNet50实现
• 添加GNN(图神经网络)层
• 支持从TensorFlow/Jax/PyTorch/Hugging Face导入模型：
• 支持黑盒推理和进一步微调
• 提供一些标准模型如ResNet50、ViT、BERT、Chinchilla等
• 重点支持从Hugging Face导入模型
• 扩展计算图操作和操作工具
• 实现NaN和Inf值的首次出现检测，并提供堆栈跟踪

• 模型保存/加载：
• 支持导出为TensorFlow的"SavedModel"格式
• 开发仅用于推理的轻量版本
• 支持将计算图预先编译(AOT)为不依赖XLA的库

• 实现同步镜像策略(数据并行)
• 支持多设备(GPU)设置
• 支持多租户(多主机)设置

• 完善多设备(如多GPU)支持
• 实现自动分布在多设备上的DistributedTrainer
• 新后端支持：
• 开发兼容WebAssembly的后端(如WebGL)，以支持浏览器中运行和训练
• 开发更快/更轻的CPU后端(可能使用ggml)

• 虽然没有具体列出，但文档提到API仍有改进空间，预计在早期阶段会有一些变化

这些计划涵盖了从模型功能扩展到基础设施改进，以及分布式和新平台支持等多个方面，旨在使gomlx成为一个更全面、强大的机器学习框架。

四. 小结

通过上面的介绍，我们看到： GoMLX站在了TensorFlow/PyTorch/Jax等机器学习框架巨头的肩膀之上 ，借鉴了他们的设计与优势，并充分复用已有的工具和库，结合Go语言的特点，但又不局限在Go之中。

GoMLX的作者就是机器学习研究员，十分清楚机器学习框架用户的诉求，提供了GoNB(A Modern Go Kernel for Jupyter Notebooks)，让用户可以在熟悉的Jupyter Notebooks中完成相关实验，并内置了图形化展示工具。GoMLX的教程和示例非常细致，并且也都以ipynb格式提供，大幅降低GoMLX上手门槛。

不过，尽管GoMLX展示了强大的潜力，但目前它仍处于初始阶段，许多核心功能尚未完全实现。但总体而言，GoMLX代表了Go语言在机器学习领域的一次重要尝试，其设计理念和目标值得关注。尽管目前还存在诸多挑战和未完成的工作，但我们可以对其未来的发展保持谨慎乐观的态度。

如果你有机器学习开发经验，不妨参与到GoMLX的开发中，为这一项目贡献力量，共同推动Go在机器学习领域的进步。

参考资料