nx_pose_backprop_gate.nx source
↩ module page · 99 lines · 5675 B
1// nx_pose_backprop_gate.nx -- proves the MULTI-LAYER pose-net backprop COMPOSES correctly (the "scale build-2"
2// de-risk): a 2-layer conv net conv1(1->2) -> ReLU -> conv2(2->1), then the loss gradient is backpropagated through
3// conv2 (nx_f32_conv2d_backward) -> ReLU backward (f32_relu_bwd) -> conv1 (nx_f32_conv2d_backward). Inputs/weights
4// are integers chosen so every pre-activation is POSITIVE (ReLU = identity) -> the net is locally linear -> the
5// chained gradients are exact small integers, verified BIT-EXACT (no tolerance). This shows the training loop
6// already gated for one layer (nx_f32_conv_train_gate) chains correctly through DEPTH -- so a real backbone+head
7// pose net trains; only SCALE + labeled data remain. 1x1 convs at H=W=1 keep the arithmetic hand-checkable.
8// expect_exit: 0
9import "nx_syscalls.nx"
10import "nx_f32_cvt.nx"
11import "nx_f32_conv2d.nx"
12import "nx_f32_conv2d_backward.nx"
13import "nx_pose_cnn.nx"
14import "nx_f32_train_ops.nx"
15
16func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return sys_write(1, s, n) }
17func gn(v: i64) -> i64 {
18 let bb: *u8 = sys_mmap(28); var m: i64 = v
19 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m }
20 let t: *u8 = sys_mmap(28); var k: i64 = 0
21 if m == 0 { t[0] = 48 as u8; k = 1 }
22 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
23 var i: i64 = 0; while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
24 return sys_write(1, bb, k)
25}
26
27func main(argc: i64, argv: *i64) -> i64 {
28 var pass: i64 = 0
29 let inp: *i64 = sys_mmap(8) as *i64 // [C_in=1,1,1]
30 let w1: *i64 = sys_mmap(8 * 2) as *i64 // [C_out=2,C_in=1,1,1]
31 let b1: *i64 = sys_mmap(8 * 2) as *i64
32 let a1: *i64 = sys_mmap(8 * 2) as *i64 // hidden (post-ReLU) [2,1,1]
33 let w2: *i64 = sys_mmap(8 * 2) as *i64 // [C_out=1,C_in=2,1,1]
34 let b2: *i64 = sys_mmap(8) as *i64
35 let pred:*i64 = sys_mmap(8) as *i64
36 let tgt: *i64 = sys_mmap(8) as *i64
37 let dpred:*i64 = sys_mmap(8) as *i64
38 let da1: *i64 = sys_mmap(8 * 2) as *i64 // = dinput of conv2
39 let dz1: *i64 = sys_mmap(8 * 2) as *i64 // after relu backward
40 let dw2: *i64 = sys_mmap(8 * 2) as *i64
41 let db2: *i64 = sys_mmap(8) as *i64
42 let dw1: *i64 = sys_mmap(8 * 2) as *i64
43 let db1: *i64 = sys_mmap(8 * 2) as *i64
44 let dx: *i64 = sys_mmap(8) as *i64 // = dinput of conv1
45
46 // x=2; w1=[3,5] b1=[1,2] -> z1=[7,12] (>0, ReLU active) ; w2=[2,1] b2=[3] -> pred = 2*7+1*12+3 = 29 ; tgt=20 -> e=9
47 inp[0] = nx_i32_to_f32(2)
48 w1[0] = nx_i32_to_f32(3); w1[1] = nx_i32_to_f32(5); b1[0] = nx_i32_to_f32(1); b1[1] = nx_i32_to_f32(2)
49 w2[0] = nx_i32_to_f32(2); w2[1] = nx_i32_to_f32(1); b2[0] = nx_i32_to_f32(3)
50 tgt[0] = nx_i32_to_f32(20)
51
52 // forward: conv1 -> a1 ; ReLU(a1) ; conv2 -> pred
53 nx_f32_conv2d_forward(inp, 1, 1, 1, 1, w1, 2, 1, 1, 1, 0, b1, a1)
54 pose_cnn_relu(a1, 2)
55 nx_f32_conv2d_forward(a1, 1, 2, 1, 1, w2, 1, 1, 1, 1, 0, b2, pred)
56
57 // P0 sanity: pred == 29
58 if pred[0] == nx_i32_to_f32(29) { pass = pass + 1; gp("P0 forward pred=29 OK\n" as *u8) }
59 if pred[0] != nx_i32_to_f32(29) { gp("P0 FAIL pred=" as *u8); gn(pred[0]); gp(" want=" as *u8); gn(nx_i32_to_f32(29)); gp("\n" as *u8) }
60
61 // loss grad: dpred = pred - tgt = 9
62 f32_mse_grad(pred, tgt, dpred, 1)
63
64 // backward conv2: dw2 = e*a1 = [63,108], db2 = 9, dinput2 = da1 = e*w2 = [18,9]
65 nx_f32_conv2d_backward(a1, 1, 2, 1, 1, w2, 1, 1, 1, 1, 0, dpred, da1, dw2, db2)
66 if dw2[0] == nx_i32_to_f32(63) { if dw2[1] == nx_i32_to_f32(108) { if db2[0] == nx_i32_to_f32(9) {
67 pass = pass + 1; gp("P1 conv2 grads dw2=[63,108] db2=9 OK\n" as *u8)
68 } } }
69 if dw2[0] != nx_i32_to_f32(63) { gp("P1 FAIL dw2=[" as *u8); gn(dw2[0]); gp("," as *u8); gn(dw2[1]); gp("]\n" as *u8) }
70
71 // P2 the propagated hidden grad da1 = [18,9]
72 if da1[0] == nx_i32_to_f32(18) { if da1[1] == nx_i32_to_f32(9) {
73 pass = pass + 1; gp("P2 da1 (dinput of conv2) = [18,9] OK\n" as *u8)
74 } }
75 if da1[0] != nx_i32_to_f32(18) { gp("P2 FAIL da1=[" as *u8); gn(da1[0]); gp("," as *u8); gn(da1[1]); gp("]\n" as *u8) }
76
77 // ReLU backward: a1>0 -> dz1 = da1 = [18,9]
78 f32_relu_bwd(a1, da1, dz1, 2)
79 if dz1[0] == nx_i32_to_f32(18) { if dz1[1] == nx_i32_to_f32(9) {
80 pass = pass + 1; gp("P3 relu backward (active) dz1=[18,9] OK\n" as *u8)
81 } }
82 if dz1[0] != nx_i32_to_f32(18) { gp("P3 FAIL dz1=[" as *u8); gn(dz1[0]); gp("," as *u8); gn(dz1[1]); gp("]\n" as *u8) }
83
84 // backward conv1: dw1 = dz1*x = [36,18], db1 = dz1 = [18,9], dinput1 = dx = sum dz1*w1 = 18*3+9*5 = 99
85 nx_f32_conv2d_backward(inp, 1, 1, 1, 1, w1, 2, 1, 1, 1, 0, dz1, dx, dw1, db1)
86 if dw1[0] == nx_i32_to_f32(36) { if dw1[1] == nx_i32_to_f32(18) { if db1[0] == nx_i32_to_f32(18) { if db1[1] == nx_i32_to_f32(9) {
87 pass = pass + 1; gp("P4 conv1 grads dw1=[36,18] db1=[18,9] OK\n" as *u8)
88 } } } }
89 if dw1[0] != nx_i32_to_f32(36) { gp("P4 FAIL dw1=[" as *u8); gn(dw1[0]); gp("," as *u8); gn(dw1[1]); gp("]\n" as *u8) }
90
91 // P5 the input gradient dx = 99 (chained all the way back through both layers + ReLU)
92 if dx[0] == nx_i32_to_f32(99) { pass = pass + 1; gp("P5 dx (input grad, full chain) = 99 OK\n" as *u8) }
93 if dx[0] != nx_i32_to_f32(99) { gp("P5 FAIL dx=" as *u8); gn(dx[0]); gp(" want=" as *u8); gn(nx_i32_to_f32(99)); gp("\n" as *u8) }
94
95 gp("POSE-BACKPROP-GATE pass=" as *u8); gn(pass); gp("/6\n" as *u8)
96 if pass == 6 { gp("POSE-BACKPROP-GATE GREEN 6/6 (conv->ReLU->conv backprop composes EXACT) -- multi-layer pose net trains\n" as *u8); sys_exit(0) }
97 sys_exit(1)
98 return 0
99}