第9章：训练管线——Dataset、Trainer 与 Loop

9.1 第 8 章你写了什么

构建训练图。遍历样本。累加梯度。手动改参数。打印 loss。

60 行。像写汇编。

GoMLX 提供了一套高层 API，把这 60 行压成 15 行。但你不会迷失——因为你知道里面在干什么。

9.2 Dataset：数据从哪来

train.Dataset 只有一个核心方法：

type Dataset interface {
    Name() string
    Reset()
    Yield() (spec any, inputs, labels []*tensors.Tensor, err error)
}

Trainer 每次要数据时，就调用 Yield()。inputs 和 labels 都是 []*tensors.Tensor 切片——即使只有一张输入、一个标签，也得包成切片。

对于第 8 章那 100 条内存数据，最快的办法是用 datasets.InMemoryFromData：

import "github.com/gomlx/gomlx/pkg/ml/datasets"

mds, err := datasets.InMemoryFromData(backend, "linear",
    []any{xData},   // inputs
    []any{yData})  // labels
if err != nil {
    log.Fatal(err)
}

InMemoryFromData 把数据一次性读进内存，放在 backend 对应的设备上。它返回的 *InMemoryDataset 可以链式配置：

ds := mds.BatchSize(10, false).Infinite(true)

• BatchSize(10, false)：每步 yield 10 条。最后一个 batch 不足 10 条也保留。
• Infinite(true)：数据跑完一轮自动从头来。

如果你的数据在文件里、数据库里，就自己实现 train.Dataset 接口。接口只有三个方法，不复杂。

9.3 Trainer：一步训练的执行者

Trainer 把"前向 → 损失 → 梯度 → 更新"这四步包成一步。你需要给它四样东西：

1. 一个 *context.Context：存所有变量和超参数。
2. 一个 ModelFn：给定输入，建图算预测。
3. 一个 LossFn：给定预测和标签，算损失。
4. 一个 Optimizer：怎么根据梯度更新变量。

ctx := context.New()
ctx.SetParam(optimizers.ParamLearningRate, 0.01)

modelFn := func(ctx *context.Context, spec any, inputs []*graph.Node) (predictions []*graph.Node) {
    x := inputs[0]
    pred := layers.DenseWithBias(ctx.In("linear"), x, 1)
    return []*graph.Node{pred}
}

trainer := train.NewTrainer(backend, ctx, modelFn,
    losses.MeanSquaredError,
    optimizers.StochasticGradientDescent(),
    nil, nil)

losses.MeanSquaredError 接收两个 []*graph.Node 切片，不是单个 *Node。这是为了支持多输出的模型。对于单输出，直接传就行。

Trainer 对外暴露的核心方法：TrainStep(spec, inputs, labels) 跑一步训练；EvalStep(spec, inputs, labels) 跑一步评估；Eval(dataset) 在整个 dataset 上评估。

9.4 Loop：把循环跑起来

Loop 管两件事：迭代 Dataset，调用 Trainer，以及触发钩子。

import "github.com/gomlx/gomlx/ui/commandline"

loop := train.NewLoop(trainer)
commandline.AttachProgressBar(loop)

metrics, err := loop.RunSteps(ds, 2000)
if err != nil {
    log.Fatal(err)
}

commandline.AttachProgressBar 在命令行画一个进度条，顺带打印 loss。RunSteps 返回 (metrics []*tensors.Tensor, err error)。如果 dataset 提前返回 io.EOF，RunSteps 不会自动续跑。所以 dataset 必须能持续供应数据（用 .Infinite(true)）。

9.5 完整代码

package main

import (
    "fmt"
    "log"

    "github.com/gomlx/gomlx/backends"
    "github.com/gomlx/gomlx/pkg/core/graph"
    "github.com/gomlx/gomlx/pkg/ml/context"
    "github.com/gomlx/gomlx/pkg/ml/datasets"
    "github.com/gomlx/gomlx/pkg/ml/layers"
    "github.com/gomlx/gomlx/pkg/ml/train"
    "github.com/gomlx/gomlx/pkg/ml/train/losses"
    "github.com/gomlx/gomlx/pkg/ml/train/optimizers"
    "github.com/gomlx/gomlx/ui/commandline"
)

func main() {
    backend, err := backends.New()
    if err != nil { log.Fatal(err) }

    // 生成数据：y = 3x + 5 + 噪声
    numSamples := 100
    xData := make([][]float64, numSamples)
    yData := make([][]float64, numSamples)
    for i := 0; i < numSamples; i++ {
        x := float64(i + 1)
        xData[i] = []float64{x}
        yData[i] = []float64{3.0*x + 5.0 + float64(i%5-2)*0.5}
    }

    mds, err := datasets.InMemoryFromData(backend, "linear", []any{xData}, []any{yData})
    if err != nil { log.Fatal(err) }
    ds := mds.BatchSize(10, false).Infinite(true)

    ctx := context.New()
    ctx.SetParam(optimizers.ParamLearningRate, 0.01)

    modelFn := func(ctx *context.Context, spec any, inputs []*graph.Node) (predictions []*graph.Node) {
        x := inputs[0]
        pred := layers.DenseWithBias(ctx.In("linear"), x, 1)
        return []*graph.Node{pred}
    }

    trainer := train.NewTrainer(backend, ctx, modelFn,
        losses.MeanSquaredError,
        optimizers.StochasticGradientDescent(),
        nil, nil)

    loop := train.NewLoop(trainer)
    commandline.AttachProgressBar(loop)

    metrics, err := loop.RunSteps(ds, 2000)
    if err != nil { log.Fatal(err) }
    if len(metrics) > 0 {
        fmt.Printf("Final loss: %v\n", metrics[0].Value())
    }
}

第 8 章的 60 行手工循环，现在变成了不到 30 行——而且你不用管梯度累加、参数更新、图编译这些琐事。但这套 API 不是黑箱。你知道 Trainer 内部做了什么。区别只是现在由框架代劳。

9.6 超参数的层次化管理

context.Context 内部是一棵作用域树。ctx.In("layer1") 进入子作用域。每个作用域可以独立设参数。

ctx := context.New()
ctx.SetParam(optimizers.ParamLearningRate, 0.01)
ctx.In("layer2").SetParam(optimizers.ParamLearningRate, 0.001)

优化器在更新变量时，会先从变量所在的作用域查学习率；查不到，就一层层往上走，直到根。这让你给不同层配不同学习率。

9.7 收束

Dataset（数据源）
    │  Yield() → 返回 (spec, inputs, labels, err)
    ▼
Loop（训练循环）
    │  每步调用 Trainer.TrainStep()
    │  每 N 步触发钩子（打印、保存、评估）
    ▼
Trainer（训练步执行器）
    │  前向：modelFn(ctx, spec, inputs) → predictions
    │  损失：lossFn(labels, predictions) → loss scalar
    │  梯度：Gradient(loss, variables...) → grads
    │  更新：optimizer.Apply(variables, grads)
    ▼
Context（变量与超参数容器）
    │  存所有权重、偏置
    │  存学习率、衰减率等超参数

第 10 章，在这套管线上搭第一个真正的神经网络。

作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》
上一章：第8章 线性回归从头写
下一章预告：第10章 神经网络——层、激活函数与前向传播