《GoMLX从入门到精通》第13章：模型保存、导出与 ONNX 互操作

Chapter 13 — Version 2

---

13.1 训练只是手段，部署才是目的

你花几个小时训练出来的模型，不能只活在 main.go 的训练代码里。你得把它变成一个可以分发的资产：

• 给运维团队扔一个二进制，他们就能启动推理服务
• 给前端同事一个 ONNX 文件，他们加载到自己的推理引擎里
• 给未来的自己一个 Checkpoint，两周后不需要重新训练就能继续调优

---

13.2 路径一：Checkpoint（原生格式）

GoMLX 最原生的持久化方式是 checkpoints.Handler——它管理 context.Context 的变量和超参数的保存与加载。

import (
    "github.com/gomlx/gomlx/pkg/ml/context"
    "github.com/gomlx/gomlx/pkg/ml/context/checkpoints"
)

// 保存
ctx := context.New()
// ... 训练代码 ...

// 创建 Handler：指定目录，保留最近 3 个 Checkpoint
handler, err := checkpoints.Build(ctx).Dir("models/my-model").Keep(3).Done()
if err != nil {
    log.Fatal(err)
}

// 手动保存一次
if err := handler.Save(); err != nil {
    log.Fatal(err)
}

加载时，Build(...).Done() 自动从目录中加载最新的 Checkpoint：

// 在另一个程序里加载
ctx2 := context.New()
handler2, err := checkpoints.Build(ctx2).Dir("models/my-model").Done()
if err != nil {
    log.Fatal(err)
}
// 此时 ctx2 已经持有和训练时一模一样的参数
// handler2.Save() 可以继续保存新的 Checkpoint

Checkpoint 里存了什么？

• 所有 context.Variable 的值（权重和偏置）
• 所有超参数（学习率、正则化系数等，通过 ctx.SetParam 设置的）
• 模型的计算图定义——不存。图定义仍然需要你的模型代码。

loop := train.NewLoop(trainer)
commandline.AttachProgressBar(loop)
if handler != nil {
    const priority = 100
    train.EveryNSteps(loop, 100, "checkpointing", priority, handler.OnStepFn)
}
// handler.OnStepFn 实现了 train.OnStepFn 接口，直接传给 EveryNSteps

最佳实践：

• 保存"最佳模型"而非"最后一步"——跟踪验证集指标，只保存验证 loss 最低的那个 Checkpoint
• 命名加上步数和验证 loss：GoMLX 自动在文件名中包含步数和时间戳
• 用 Keep(n) 控制保留数量，避免磁盘被 Checkpoint 塞满

13.3 路径二：ONNX（跨框架标准格式）

ONNX（Open Neural Network Exchange）是跨框架的模型交换格式。PyTorch、TensorFlow、JAX 都支持导出 ONNX。

GoMLX → ONNX（导出）：截至 v0.27.3，ONNX 导出 API 仍在规划中（源码 docs/TODO.md 中标记为 idea，优先级 P1）。届时你训练好的模型可以直接导出为 .onnx 文件，被任何 ONNX 运行时加载。

ONNX → GoMLX（导入）：通过独立项目 github.com/gomlx/onnx-gomlx，你可以把 Hugging Face 上的模型（PyTorch/JAX 导出）转成 GoMLX 格式，然后在纯 Go 环境里跑推理。

// 从 ONNX 加载模型（概念示例——具体 API 以 onnx-gomlx 项目文档为准）
// onnx-gomlx 将 ONNX 计算图转成 GoMLX 的 context.Context + 模型函数
// 转换完成后，和原生 GoMLX 模型一样使用
exec := context.MustNewExec(backend, ctx, modelFn)
result := exec.MustExec1(input)

ONNX 路径的价值：

• 你可以在 Python 生态训练（PyTorch/JAX），导出 ONNX，然后在 Go 服务里做推理
• 你可以直接加载 Hugging Face 上的预训练模型（BERT、Gemma 等），不需要重新训练
• 你的 GoMLX 模型可以交给用 C++/Rust/JS 的同事——ONNX 是"通用语"

13.4 路径三：纯 Go 推理二进制（零依赖部署）

这是 GoMLX 最独特的优势——编译一个包含模型权重的独立二进制文件。

GoMLX 的 checkpoints.Handler 原生支持从 go:embed 加载：

import (
    _ "embed"
    "github.com/gomlx/gomlx/pkg/ml/context"
    "github.com/gomlx/gomlx/pkg/ml/context/checkpoints"
)

//go:embed "my_model/checkpoint.json"
var myModelJson string

//go:embed "my_model/checkpoint.bin"
var myModelBin []byte

func main() {
    backend, _ := backends.New()
    ctx := context.New()

    // 从嵌入数据加载 Checkpoint——不需要临时文件
    _, err := checkpoints.Build(ctx).
        FromEmbed(myModelJson, myModelBin).
        Done()
    if err != nil {
        log.Fatal(err)
    }

    // 创建推理 Exec
    exec := context.MustNewExec(backend, ctx, modelFn)

    // 启动 HTTP 推理服务
    http.HandleFunc("/predict", func(w http.ResponseWriter, r *http.Request) {
        input := parseInput(r)
        result := exec.MustExec1(input)
        defer result.FinalizeAll()
        writeJSON(w, result.Value())
    })
    http.ListenAndServe(":8080", nil)
}

产物：一个 30-80MB 的二进制（含 GoMLX runtime + SimpleGo 后端 + 模型权重）。部署命令：scp my-server user@host: 然后 ./my-server。

对比 Python 推理服务：

• 不需要 Python 解释器
• 不需要 pip install torch transformers ...
• 不需要 CUDA 版本匹配
• 不需要 Docker（当然你也可以用）
• 同一进程内调用，无网络延迟（对比 Python 独立推理进程 + gRPC 通信）

13.5 路径选择指南

> 作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》

[上一章：第12章 实战：CIFAR-10 图像分类全流程] | [下一章：第14章 性能调优——从 Profile 到 Kernel Fusion]