《GoMLX从入门到精通》第12章：实战：CIFAR-10 图像分类全流程

Chapter 12

---

12.1 项目全景

12.2 数据准备

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

// CIFAR-10 二进制格式解析（简化——实际项目可能从文件或预加载张量读）
func loadCIFAR10() (trainDS, testDS train.Dataset) {
    // 加载训练集：50000 × [32, 32, 3] 的 uint8 图像
    trainImages := loadBinaryImages("cifar-10-batches-bin/data_batch_*.bin")
    trainLabels := loadBinaryLabels("cifar-10-batches-bin/data_batch_*.bin")

    // 数据增强（Data Augmentation）
    augmentFn := func(images, labels *tensors.Tensor) (*tensors.Tensor, *tensors.Tensor) {
        // 1. 随机水平翻转（50% 概率）
        // 2. 随机裁剪 + padding（32×32 → 40×40 → 随机裁回 32×32）
        // 3. 归一化到 [0, 1]，再标准化（减均值/除方差）
        return augmentedImages, labels
    }

    trainDS = dataset.Custom(
        "CIFAR-10 Train",
        func() (chan *train.Batch, chan error) {
            // 每批 128 张图，shuffle，应用增强
            return batchedIterator(trainImages, trainLabels, 128, true, augmentFn)
        })

    // 测试集：不增强，只归一化
    testDS = dataset.Custom(
        "CIFAR-10 Test",
        func() (chan *train.Batch, chan error) {
            return batchedIterator(testImages, testLabels, 128, false, normalizeOnly)
        })

    return
}

数据增强的意义：每张原始图片在每次训练 epoch 中都被随机变换——模型永远不会看到"完全相同的"输入两次。这极大提升了泛化能力——等价于无限扩增了训练集。

---

12.3 模型：简化 ResNet-18

ResNet 的核心创新是残差连接（Skip Connection）。一个残差块长这样：

import (
    "github.com/gomlx/gomlx/pkg/ml/context"
    "github.com/gomlx/gomlx/pkg/core/graph"
    "github.com/gomlx/gomlx/pkg/ml/layers"
    "github.com/gomlx/gomlx/pkg/ml/layers/activations"
    "github.com/gomlx/gomlx/pkg/ml/layers/batchnorm"
)

func resnetBlock(ctx *context.Context, x *graph.Node, filters int, stride int) *graph.Node {
    residual := x  // 保存短路连接
    inputChannels := x.Shape().Dimensions[len(x.Shape().Dimensions)-1]

    // 主路径：Conv → BN → ReLU → Conv → BN
    x = layers.Convolution(ctx.In("conv_a"), x).Channels(filters).
        KernelSizePerAxis(3, 3).Strides(stride).PadSame().Done()
    x = batchnorm.New(ctx.In("bn_a"), x, -1).Done()
    x = activations.Relu(x)

    x = layers.Convolution(ctx.In("conv_b"), x).Channels(filters).
        KernelSizePerAxis(3, 3).Strides(1).PadSame().Done()
    x = batchnorm.New(ctx.In("bn_b"), x, -1).Done()

    // 如果 stride > 1 或通道数变了，短路也需要投影
    if stride > 1 || inputChannels != filters {
        residual = layers.Convolution(ctx.In("proj"), residual).Channels(filters).
            KernelSize(1).Strides(stride).Done()
    }

    // 残差连接：主路径 + 短路
    x = graph.Add(x, residual)
    return activations.Relu(x)
}

残差连接为什么有效：传统深层网络中，梯度在反向传播时一层一层衰减（梯度消失）。残差连接给梯度提供了一条"高速公路"——梯度可以直接从输出跳到输入，绕过中间层。这使得 152 层甚至更深的网络也能有效训练。

完整的 ResNet-18 结构：

输入 [batch, 32, 32, 3]
  → Conv 3×3, 64, stride=1
  → ResBlock × 2, 64 通道
  → ResBlock × 2, 128 通道 (stride=2, 尺寸 16×16)
  → ResBlock × 2, 256 通道 (stride=2, 尺寸 8×8)
  → ResBlock × 2, 512 通道 (stride=2, 尺寸 4×4)
  → Global Average Pooling
  → Dense(512 → 10)
  → 输出 [batch, 10]

---

12.4 训练策略

import (
    "github.com/gomlx/gomlx/backends"
    "github.com/gomlx/gomlx/pkg/ml/context"
    "github.com/gomlx/gomlx/pkg/ml/context/checkpoints"
    "github.com/gomlx/gomlx/pkg/ml/layers/regularizers"
    "github.com/gomlx/gomlx/pkg/ml/train"
    "github.com/gomlx/gomlx/pkg/ml/train/losses"
    "github.com/gomlx/gomlx/pkg/ml/train/metrics"
    "github.com/gomlx/gomlx/pkg/ml/train/optimizers"
    "github.com/gomlx/gomlx/pkg/ml/train/commandline"
)

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

    // 1. 加载数据
    trainDS, testDS := loadCIFAR10()

    // 2. 创建 Context 和超参数
    ctx := context.New()
    ctx.SetParam(optimizers.ParamLearningRate, 0.1) // 初始学习率
    ctx.SetParam(regularizers.ParamL2, 1e-4)       // L2 正则化

    // 3. 创建 Checkpoint Handler
    checkpoint := checkpoints.Build(ctx).Dir("cifar10-checkpoint").MustDone()

    // 4. 创建 Trainer
    // 注意：SGD 默认带 1/sqrt(step) 学习率衰减。如需 Cosine Annealing，
    // 需在模型图内调用 cosineschedule.New(ctx, g, dtype).FromContext().Done()
    trainer := train.NewTrainer(backend, ctx,
        ResNet18,                         // 模型函数
        losses.SparseCategoricalCrossEntropyLogits, // 损失
        optimizers.StochasticGradientDescent(), // SGD
        []metrics.Interface{              // 训练评估指标
            metrics.NewSparseCategoricalAccuracy("accuracy", "acc"),
        },
        nil)

    // 6. 创建 Loop + 钩子
    loop := train.NewLoop(trainer)
    commandline.AttachProgressBar(loop)

    // 每 1000 步评估测试集
    evalHook := train.EveryNSteps(loop, 1000, "test_eval", 0,
        func(loop *train.Loop, metrics []*tensors.Tensor) error {
            return loop.Trainer.Eval(testDS)
        })
    loop.AttachHook(evalHook)

    // 每 5000 步保存 Checkpoint
    saveHook := train.EveryNSteps(loop, 5000, "save", 0,
        func(loop *train.Loop, _ []*tensors.Tensor) error {
            return checkpoint.Save()
        })
    loop.AttachHook(saveHook)

    // 7. 跑训练
    trainMetrics, err := loop.RunSteps(trainDS, 64000)
    if err != nil {
        log.Fatal(err)
    }
    _ = trainMetrics
    fmt.Println("Training completed.")
}

为什么 Cosine Decay？ 训练开始时大步子（lr=0.1），快速靠近最优区。接近结束时小步精细调整。Cosine 曲线比手动阶梯式衰减更平滑，在实践中往往给出更好的最终结果。

---

12.5 评估与分析

训练完成后，不只看一个数字：

// 1. 在测试集上的最终准确率
testMetrics, _ := trainer.Eval(testDS)
fmt.Printf("Test accuracy: %.2f%%\n", testMetrics[0].Value())

// 2. 混淆矩阵——哪些类别容易搞混？
// （需要你在 testDS 上逐批推理，对比预测和真实标签）
// 典型发现：猫 vs 狗是 CIFAR-10 最难区分的两个类

// 3. 错误案例分析
// 做完混淆矩阵后，抽出"模型判错"的几张图片肉眼观察。
// 常见原因：
//   - 图像太模糊（32×32 本身分辨率低）
//   - 物体被遮挡或只在角落
//   - 标注本身有歧义（那张图里的到底是"鹿"还是"马"）

---

12.6 模型导出

训练好的模型，不只是"在训练代码里能用"。要把它变成一个独立的推理服务。

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

// 1. 保存 Checkpoint
checkpoint := checkpoints.Build(ctx).Dir("cifar10-resnet18").MustDone()
checkpoint.Save()

// 2. 在推理程序中加载
func loadModel(backend backends.Backend, ckptPath string) *context.Context {
    ctx := context.New()
    checkpoints.Load(ctx).Dir(ckptPath).MustDone() // 加载所有权重
    return ctx
}

// 3. 创建纯推理 Exec
inferenceExec, err := context.NewExec(backend, ctx, ResNet18)
if err != nil {
    log.Fatal(err)
}

// 4. 推理——只需要输入图像，不需要训练数据
image := loadAndPreprocessImage("cat.jpg")
logits := inferenceExec.MustExec1(image)
predictedClass := argmax(logits.Value())

---

12.7 收束：从知道到做到

第 8 章你手写了训练循环，知道了里面每一根水管怎么接。

第 9 章你拿到了标准件——Dataset、Trainer、Loop。

第 10 章你学会了搭积木，从全连接层到 Dropout、BatchNorm。

第 11 章你掌握了 CNN 和 Attention 两种最重要的架构。

第 12 章你把这些全部组装在一起，跑完了一个完整的图像分类项目。Code 是模板，知识是骨架。从这一章出发，你可以修改 ResNet-18 的深度、换用不同的数据增强策略、接入你自己的数据集。

第 13 章讲模型导出和 ONNX 互通——让你的模型从"项目内的函数"变成"可以分发、部署、商用的资产"。

---

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

[上一章：第11章 卷积与注意力——两大基础架构] | [下一章：第13章 模型保存、导出与 ONNX 互操作]