同一个小模型训练到收敛

class Cnn(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv = nn.Conv1d(1, CH, K)
    self.fc = nn.Linear(FLAT, CLASSES)

  def forward(self, x):
    return self.fc(torch.relu(self.conv(x)).flatten(1))

#[derive(Module, Debug)]
struct Cnn {
  conv: Conv1d,
  fc: Linear,
  relu: Relu,
}

impl Cnn {
  fn forward(&self, x: Tensor<3>) -> Tensor<2> {
    self.fc.forward(self.relu.forward(self.conv.forward(x)).flatten(1, 2))
  }
}

python -I train_both.py > pt.json
cargo run --release --offline > burn.json
python -I compare.py pt.json burn.json

sgd.losses               n=100  max|diff|=1.79e-07
sgd.final.loss           n=  1  max|diff|=1.49e-08
sgd.final.correct        n=  1  max|diff|=0.00e+00
sgd.params.conv.weight   n= 20  max|diff|=7.15e-07
sgd.params.conv.bias     n=  4  max|diff|=5.96e-07
sgd.params.fc.weight     n=144  max|diff|=4.77e-07
sgd.params.fc.bias       n=  3  max|diff|=4.92e-07
adam.losses              n=100  max|diff|=2.77e-06
adam.final.loss          n=  1  max|diff|=7.75e-07
adam.final.correct       n=  1  max|diff|=0.00e+00
adam.params.conv.weight  n= 20  max|diff|=3.04e-06
adam.params.conv.bias    n=  4  max|diff|=2.62e-06
adam.params.fc.weight    n=144  max|diff|=3.52e-06
adam.params.fc.bias      n=  3  max|diff|=1.61e-06
worst=3.52e-06 tol=1e-06 MISMATCH

step  torch loss  burn loss   |dloss|   max|dparam|  torch-f64  burn-f64   f64-f64
   1  1.13687289  1.13687277  1.19e-07  2.98e-08     2.40e-08   1.93e-08  5.55e-17
   2  1.10447657  1.10447657  0.00e+00  2.98e-08     2.84e-08   2.37e-08  5.55e-17
   5  1.06677926  1.06677914  1.19e-07  4.47e-08     3.09e-08   2.89e-08  8.33e-17
  10  0.98627138  0.98627144  5.96e-08  4.47e-08     6.23e-08   3.99e-08  1.39e-16
  20  0.84775448  0.84775442  5.96e-08  8.94e-08     1.12e-07   1.16e-07  8.88e-16
  30  0.64109868  0.64109868  0.00e+00  1.79e-07     2.01e-07   1.63e-07  2.66e-15
  40  0.50108528  0.50108522  5.96e-08  2.98e-07     4.37e-07   2.00e-07  6.33e-15
  50  0.38287219  0.38287210  8.94e-08  4.77e-07     7.14e-07   2.49e-07  1.17e-14
  60  0.22960071  0.22960062  8.94e-08  5.36e-07     8.23e-07   3.72e-07  1.38e-14
  70  0.15450853  0.15450841  1.19e-07  4.77e-07     7.37e-07   3.26e-07  1.53e-14
  80  0.10406791  0.10406786  4.47e-08  4.77e-07     8.16e-07   4.42e-07  1.88e-14
  90  0.07615292  0.07615286  5.96e-08  7.15e-07     9.32e-07   4.99e-07  2.15e-14
 100  0.06210175  0.06210171  3.73e-08  7.15e-07     1.05e-06   5.75e-07  2.32e-14

step  torch loss  burn loss   |dloss|   max|dparam|  torch-f64  burn-f64   f64-f64
   1  1.13687289  1.13687277  1.19e-07  4.69e-07     4.91e-07   2.15e-08  2.23e-10
   2  1.11771631  1.11771631  0.00e+00  5.07e-07     5.05e-07   2.38e-08  3.44e-08
   5  1.08280766  1.08280778  1.19e-07  4.21e-07     4.82e-07   3.74e-08  1.58e-07
  10  1.03940225  1.03940248  2.38e-07  3.87e-07     4.66e-07   5.62e-08  3.91e-07
  20  0.97808975  0.97809017  4.17e-07  7.60e-07     4.95e-07   8.15e-08  7.85e-07
  30  0.83425587  0.83425671  8.34e-07  1.19e-06     4.84e-07   1.19e-07  1.23e-06
  40  0.69612402  0.69612581  1.79e-06  1.65e-06     4.64e-07   1.06e-07  1.69e-06
  50  0.58540988  0.58541220  2.32e-06  2.32e-06     4.48e-07   1.92e-07  2.35e-06
  60  0.38402307  0.38402560  2.53e-06  2.92e-06     4.61e-07   1.76e-07  2.93e-06
  70  0.24789160  0.24789384  2.24e-06  3.46e-06     4.66e-07   2.21e-07  3.36e-06
  80  0.17158426  0.17158614  1.88e-06  3.46e-06     4.36e-07   3.06e-07  3.66e-06
  90  0.10003939  0.10004054  1.15e-06  3.52e-06     4.45e-07   3.25e-07  3.70e-06
 100  0.07264095  0.07264175  7.97e-07  3.52e-06     4.80e-07   3.55e-07  3.71e-06

  t  burn f32     f64          rel err
  1  0.316225648  0.316225655  -2.23e-08
  2  0.235314354  0.235315158  -3.42e-06
  5  0.172497287  0.172497706  -2.43e-06
 10  0.153187588  0.153188071  -3.15e-06
 20  0.160230935  0.160231541  -3.79e-06
 30  0.179566771  0.179567461  -3.84e-06
 40  0.201034531  0.201035340  -4.02e-06
 50  0.222036406  0.222037423  -4.58e-06
 60  0.241810858  0.241811856  -4.12e-06
 70  0.260235041  0.260236113  -4.12e-06
 80  0.277403384  0.277404527  -4.12e-06
 90  0.293464452  0.293465549  -3.74e-06
100  0.308562845  0.308564004  -3.76e-06

c, e = self.fac[self.t]  # bin factors 打印的 combined_factor 和 epsilon * bias_correction2_sqrt
m = m * self.b1 + g * (1 - self.b1)
v = v * self.b2 + g * g * (1 - self.b2)
p -= (m * c / (v.sqrt() + e)) * self.lr

step  |dloss|   max|dparam|  where
   1  0.00e+00  7.63e-17     fc.weight[67]
  10  0.00e+00  8.33e-17     fc.weight[106]
  50  1.11e-16  2.36e-16     conv.weight[5]
 100  4.16e-17  6.94e-16     fc.weight[0]

adam.losses              n=100  max|diff|=4.25e-04
adam.final.loss          n=  1  max|diff|=1.58e-05
adam.final.correct       n=  1  max|diff|=0.00e+00
adam.params.conv.weight  n= 20  max|diff|=5.67e-03
adam.params.conv.bias    n=  4  max|diff|=3.16e-03
adam.params.fc.weight    n=144  max|diff|=4.57e-03
adam.params.fc.bias      n=  3  max|diff|=4.78e-04
worst=5.67e-03 tol=1e-05 MISMATCH

# 一维小 CNN 在合成数据上训练 100 步（SGD 和 Adam 各一次），输出每步 loss 和最后的参数
import json, math, sys
import torch
from torch import nn

N, LEN, CLASSES, BATCH, STEPS = 96, 16, 3, 16, 100
CH, K = 4, 5
FLAT = CH * (LEN - K + 1)  # 4 * 12 = 48

def data():
  # 先在 f64 里算，最后舍入成 f32 一次；burn 侧同样的公式、同样的顺序
  xs, ys = [], []
  for k in range(N):
    c = k % CLASSES
    f = 0.3 * (c + 1)
    xs.append([[math.sin(f * t + 0.37 * k) + 0.3 * math.sin(2.3 * t + 1.1 * k) for t in range(LEN)]])
    ys.append(c)
  return torch.tensor(xs, dtype=torch.float32), torch.tensor(ys)

class Cnn(nn.Module):
  def __init__(self):
    super().__init__()
    self.conv = nn.Conv1d(1, CH, K)
    self.fc = nn.Linear(FLAT, CLASSES)

  def forward(self, x):
    return self.fc(torch.relu(self.conv(x)).flatten(1))

def build():
  m = Cnn()
  with torch.no_grad():
    # conv.weight [out, in, k]，fc.weight [out, in]：按下标 i 生成，两边相同
    m.conv.weight.copy_(torch.tensor([0.4 * math.sin(i + 1) for i in range(CH * K)]).reshape(CH, 1, K))
    m.conv.bias.copy_(torch.tensor([0.02 * o - 0.03 for o in range(CH)]))
    m.fc.weight.copy_(torch.tensor([0.15 * math.cos(i + 1) for i in range(CLASSES * FLAT)]).reshape(CLASSES, FLAT))
    m.fc.bias.zero_()
  return m

def train(make_opt):
  x, y = data()
  model = build()
  opt = make_opt(model.parameters())
  losses = []
  for step in range(STEPS):
    i = (step % (N // BATCH)) * BATCH  # 不打乱，6 个批轮流用
    loss = nn.functional.cross_entropy(model(x[i:i + BATCH]), y[i:i + BATCH])
    losses.append(loss.item())
    opt.zero_grad()
    loss.backward()
    opt.step()
  with torch.no_grad():
    out = model(x)
    final = {
      "loss": nn.functional.cross_entropy(out, y).item(),
      "correct": (out.argmax(1) == y).sum().item(),
    }
  params = {n: p.tolist() for n, p in model.named_parameters()}
  return {"losses": losses, "final": final, "params": params}

json.dump({
  "sgd": train(lambda p: torch.optim.SGD(p, lr=0.5)),
  "adam": train(lambda p: torch.optim.Adam(p, lr=0.01, eps=1e-8)),
}, sys.stdout)

[package]
name = "handson-train"
version = "0.1.0"
edition = "2024"

[dependencies]
burn = { version = "=0.22.1", default-features = false, features = ["std", "flex", "autodiff", "optim"] }
serde_json = "1"

//! 和 train_both.py 同一个一维小 CNN、同样的权重和数据，SGD 和 Adam 各训练 100 步，
//! 输出同样结构的 JSON。fc.weight 输出时转回 PyTorch 的 [out, in]。

use burn::module::{Module, Param};
use burn::nn::conv::{Conv1d, Conv1dConfig};
use burn::nn::loss::CrossEntropyLossConfig;
use burn::nn::{Linear, LinearConfig, Relu};
use burn::optim::{AdamConfig, GradientsParams, ModuleOptimizer, SgdConfig};
use burn::tensor::{Device, Int, Tensor, TensorData, s};
use serde_json::{Value, json};

const N: usize = 96;
const LEN: usize = 16;
const CLASSES: usize = 3;
const BATCH: usize = 16;
const STEPS: usize = 100;
const CH: usize = 4;
const K: usize = 5;
const FLAT: usize = CH * (LEN - K + 1); // 4 * 12 = 48

#[derive(Module, Debug)]
struct Cnn {
  conv: Conv1d,
  fc: Linear,
  relu: Relu,
}

impl Cnn {
  fn forward(&self, x: Tensor<3>) -> Tensor<2> {
    self.fc.forward(self.relu.forward(self.conv.forward(x)).flatten(1, 2))
  }
}

/// 先在 f64 里算，最后舍入成 f32 一次，和 Python 逐位相同
fn floats<const D: usize>(v: Vec<f64>, shape: [usize; D], device: &Device) -> Tensor<D> {
  let v: Vec<f32> = v.into_iter().map(|a| a as f32).collect();
  Tensor::from_data(TensorData::new(v, shape), device)
}

fn data(device: &Device) -> (Tensor<3>, Tensor<1, Int>) {
  let mut xs = Vec::with_capacity(N * LEN);
  let mut ys = Vec::with_capacity(N);
  for k in 0..N {
    let c = k % CLASSES;
    let f = 0.3 * (c + 1) as f64;
    for t in 0..LEN {
      let (k, t) = (k as f64, t as f64);
      xs.push((f * t + 0.37 * k).sin() + 0.3 * (2.3 * t + 1.1 * k).sin());
    }
    ys.push(c as i64);
  }
  let y = Tensor::from_data(TensorData::new(ys, [N]), device);
  (floats(xs, [N, 1, LEN], device), y)
}

fn build(device: &Device) -> Cnn {
  let mut m = Cnn {
    conv: Conv1dConfig::new(1, CH, K).init(device),
    fc: LinearConfig::new(FLAT, CLASSES).init(device),
    relu: Relu::new(),
  };
  // conv.weight 和 PyTorch 一样是 [out, in, k]；Linear 的权重是 [in, out]，按 PyTorch 的 [out, in] 生成再转置
  let w1 = (0..CH * K).map(|i| 0.4 * ((i + 1) as f64).sin()).collect();
  let b1 = (0..CH).map(|o| 0.02 * o as f64 - 0.03).collect();
  let w2 = (0..CLASSES * FLAT).map(|i| 0.15 * ((i + 1) as f64).cos()).collect();
  m.conv.weight = Param::from_tensor(floats(w1, [CH, 1, K], device));
  m.conv.bias = Some(Param::from_tensor(floats(b1, [CH], device)));
  m.fc.weight = Param::from_tensor(floats::<2>(w2, [CLASSES, FLAT], device).transpose());
  m.fc.bias = Some(Param::from_tensor(Tensor::zeros([CLASSES], device)));
  m
}

/// 转成 f64 再写 JSON（f32 转 f64 是精确的），和 Python 的 tolist() 一样
fn values<const D: usize>(t: Tensor<D>) -> Value {
  json!(t.into_data().convert::<f64>().try_to_vec::<f64>().unwrap_or_default())
}

fn params(m: &Cnn) -> Value {
  let bias = |b: &Option<Param<Tensor<1>>>| b.as_ref().map(|b| values(b.val())).unwrap_or(Value::Null);
  json!({
    "conv.weight": values(m.conv.weight.val()),
    "conv.bias": bias(&m.conv.bias),
    "fc.weight": values(m.fc.weight.val().transpose()),
    "fc.bias": bias(&m.fc.bias),
  })
}

fn train(mut optim: ModuleOptimizer, lr: f64) -> Value {
  let device = Device::flex().autodiff();
  let (x, y) = data(&device);
  let ce = CrossEntropyLossConfig::new().init(&device);
  let mut model = build(&device);
  let mut losses = Vec::with_capacity(STEPS);
  for step in 0..STEPS {
    let i = (step % (N / BATCH)) * BATCH; // 不打乱，6 个批轮流用
    let xb = x.clone().slice(s![i..i + BATCH]);
    let yb = y.clone().slice(s![i..i + BATCH]);
    let loss = ce.forward(model.forward(xb), yb);
    losses.push(loss.clone().into_scalar::<f64>());
    let grads = GradientsParams::from_grads(loss.backward(), &model);
    model = optim.step(lr, model, grads);
  }
  // 不带 autodiff 的副本和数据，相当于 model.eval() 加 torch.no_grad()
  let eval = model.valid();
  let (x, y) = data(&device.inner());
  let out = eval.forward(x);
  let loss: f64 = CrossEntropyLossConfig::new().init(&out.device()).forward(out.clone(), y.clone()).into_scalar();
  let correct: i64 = out.argmax(1).squeeze_dim::<1>(1).equal(y).int().sum().into_scalar();
  json!({
    "losses": losses,
    "final": { "loss": loss, "correct": correct },
    "params": params(&model),
  })
}

fn main() {
  let out = json!({
    "sgd": train(SgdConfig::new().init(), 0.5),
    "adam": train(AdamConfig::new().with_epsilon(1e-8).init(), 0.01),
  });
  println!("{out}");
}

torch.float32  grad 1.003785996545048e-06    update -0.009901359677314758
torch.float64  grad 1.0088567935009896e-06   update -0.009901850779581708

fc.weight[1, 19] grad Some([1.0086492920891033e-6])
configure F64: Ok(())
fc.weight[1, 19] grad Some([1.0088567935017457e-6])