第8章：你的第一个模型——线性回归从头写

8.1 问题设定：拟合一条直线

你有一堆数据点——100 个 (x, y) 对。它们大致分布在一条直线附近：

真正的规律：y = 3x + 5 + 随机噪声
你的数据：  x=[1,2,3,...,100], y≈[8,11,14,...,305]

你不知道真实斜率和截距是多少。你要让模型自己找出来。

8.2 从零到一：五步训练循环

第一步：定义模型

import (
    _ "github.com/gomlx/gomlx/backends/default"
    "github.com/gomlx/gomlx/backends"
    "github.com/gomlx/gomlx/pkg/core/dtypes"
    "github.com/gomlx/gomlx/pkg/core/graph"
    "github.com/gomlx/gomlx/pkg/core/tensors"
    "github.com/gomlx/gomlx/pkg/core/shapes"
)

func modelFn(x *graph.Node, w, b *graph.Node) *graph.Node {
    wx := graph.Mul(x, w)
    return graph.Add(wx, b)
}

第二步：定义损失函数

func lossFn(pred, target *graph.Node) *graph.Node {
    diff := graph.Sub(pred, target)
    sq := graph.Square(diff)
    return graph.ReduceAllMean(sq)
}

第三步：计算梯度

func trainStepFn(x, target, w, b *graph.Node) (loss, gradW, gradB *graph.Node) {
    pred := modelFn(x, w, b)
    loss = lossFn(pred, target)
    grads := graph.Gradient(loss, w, b)
    gradW, gradB = grads[0], grads[1]
    return
}

第四步：更新参数

参数更新不在计算图里——它在 Go 代码里：

func update(wTensor, gradWTensor *tensors.Tensor, lr float64) {
    wTensor.MaterializeLocal()
    gradWTensor.MaterializeLocal()
    w := tensors.ToScalar[float64](wTensor)
    gradW := tensors.ToScalar[float64](gradWTensor)
    newW := w - lr * gradW
    err := tensors.MutableFlatData[float64](wTensor, func(flat []float64) {
        flat[0] = newW
    })
    if err != nil {
        log.Fatalf("update w: %v", err)
    }
}

第五步：循环

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

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

    wVal := tensors.FromAnyValue(float64(0.0))
    bVal := tensors.FromAnyValue(float64(0.0))

    lr := 0.0001
    steps := 500

    for step := 0; step < steps; step++ {
        exec := graph.MustNewExec(backend,
            func(xNode, targetNode *graph.Node) (loss, gradW, gradB *graph.Node) {
                wNode := graph.Parameter(xNode.Graph(), "w", shapes.Make(dtypes.Float64))
                bNode := graph.Parameter(xNode.Graph(), "b", shapes.Make(dtypes.Float64))
                return trainStepFn(xNode, targetNode, wNode, bNode)
            })

        totalLoss := 0.0
        accumW, accumB := 0.0, 0.0
        for i := 0; i < numSamples; i++ {
            results := exec.MustExec(
                tensors.FromAnyValue(xData[i]),
                tensors.FromAnyValue(yData[i]),
                wVal, bVal,
            )
            results[0].MaterializeLocal()
            results[1].MaterializeLocal()
            results[2].MaterializeLocal()
            totalLoss += tensors.ToScalar[float64](results[0])
            accumW += tensors.ToScalar[float64](results[1])
            accumB += tensors.ToScalar[float64](results[2])
        }

        avgLoss := totalLoss / float64(numSamples)
        avgGradW := accumW / float64(numSamples)
        avgGradB := accumB / float64(numSamples)

        err := tensors.MutableFlatData[float64](wVal, func(data []float64) {
            data[0] -= lr * avgGradW
        })
        if err != nil { log.Fatalf("update w: %v", err) }
        err = tensors.MutableFlatData[float64](bVal, func(data []float64) {
            data[0] -= lr * avgGradB
        })
        if err != nil { log.Fatalf("update b: %v", err) }

        if step%50 == 0 {
            wVal.MaterializeLocal()
            bVal.MaterializeLocal()
            fmt.Printf("Step %3d: loss=%.3f, w=%.3f, b=%.3f\n",
                step, avgLoss,
                tensors.ToScalar[float64](wVal),
                tensors.ToScalar[float64](bVal))
        }
    }
}

8.3 观察：参数在"移动"

训练过程的输出大概是这样：

Step   0: loss=1008.427, w=0.000, b=0.000
Step  50: loss=342.055,  w=1.476, b=2.344
Step 100: loss=117.854,  w=2.214, b=3.575
Step 150: loss=41.093,   w=2.615, b=4.214
Step 200: loss=14.633,   w=2.818, b=4.543
Step 250: loss=5.446,    w=2.921, b=4.710
...
Step 450: loss=0.225,    w=2.999, b=5.049
Step 500: loss=0.218,    w=3.000, b=5.052

w 从 0 走到接近 3，b 从 0 走到接近 5。loss 从 1000+ 降到 0.2。模型没有"被告知"真实参数——它只是沿着梯度一步一步走过去。

8.4 这段漫长代码的意义

你刚手写了 60 行 GoMLX 代码来完成一个线性回归训练。这在任何高层 ML 框架里只需要 5 行。但你获得了几样东西：

1. 对训练本质的直接感受——每一步都在你的控制之中：前向算预测、算损失、算梯度、更新参数。没有"调用一个黑箱函数然后等结果"。

2. 对"学习率"的直观理解——把 lr 从 0.0001 改成 0.01，训练立刻爆炸（loss → NaN）；改成 1e-6，收敛极慢。你会真正理解"学习率是步子大小"。

3. 为第 9 章做的心理铺垫——现在你知道了：第 9 章的 train.Trainer 只是把这 60 行自动化了。它不是魔法，是工程。

8.5 进入高层 API 之前

你现在有了一张完整的手写训练蓝图。第 9 章，GoMLX 用 Dataset、Trainer、Loop 把这个蓝图变成了工业级代码——你不再需要手动管理变量更新、不再需要手写循环。

但你不会迷失。因为你知道里面是什么。

作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》
上一章：第7章 自动微分
下一章预告：第9章 训练管线——用高层API写出工业级训练代码