nx_f32_train_ops.nx source
↩ module page · 47 lines · 2003 B
1// nx_f32_train_ops.nx -- the small reusable software-f32 TRAINING ops that close the conv training loop, composing
2// the gated f32 tower (nx_f32_mul/add/sub/gt) with the new nx_f32_conv2d_backward. These are the shared primitives
3// the sovereign pose-net training (operator: "build 2") uses every step: MSE loss + its gradient (heatmap
4// regression is MSE over the target heatmaps), the ReLU backward (gradient flows only where the activation fired),
5// and an SGD parameter step. No optimizer state here (Adam lives in nx_f32_adam); this is the minimal, gate-able
6// core. license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_f32.nx"
9import "nx_f32_cvt.nx"
10const K_MAGIC_8192: i64 = 8192
11
12// loss = 0.5 * sum_i (pred_i - target_i)^2 (returns an f32). The standard heatmap-regression objective.
13func f32_mse_loss(pred: *i64, target: *i64, n: i64) -> i64 {
14 let half: i64 = nx_q14_to_f32(K_MAGIC_8192) // 0.5
15 var acc: i64 = 0 // +0.0
16 var i: i64 = 0
17 while i < n {
18 let d: i64 = nx_f32_sub(pred[i], target[i])
19 acc = nx_f32_add(acc, nx_f32_mul(half, nx_f32_mul(d, d)))
20 i = i + 1
21 }
22 return acc
23}
24
25// dLoss/dPred_i = pred_i - target_i -> the dOut fed into the network backward.
26func f32_mse_grad(pred: *i64, target: *i64, dpred: *i64, n: i64) -> i64 {
27 var i: i64 = 0
28 while i < n { dpred[i] = nx_f32_sub(pred[i], target[i]); i = i + 1 }
29 return 0
30}
31
32// ReLU backward: gradient passes where the (post-ReLU) activation is > 0, else 0. Composes with pose_cnn_relu.
33func f32_relu_bwd(post: *i64, dpost: *i64, dpre: *i64, n: i64) -> i64 {
34 var i: i64 = 0
35 while i < n {
36 if nx_f32_gt(post[i], 0) == 1 { dpre[i] = dpost[i] } else { dpre[i] = 0 }
37 i = i + 1
38 }
39 return 0
40}
41
42// SGD step in place: w_i <- w_i - lr * dw_i.
43func f32_sgd_step(w: *i64, dw: *i64, lr: i64, n: i64) -> i64 {
44 var i: i64 = 0
45 while i < n { w[i] = nx_f32_sub(w[i], nx_f32_mul(lr, dw[i])); i = i + 1 }
46 return 0
47}