nx_pose_student_distill.nx source
↩ module page · 156 lines · 10242 B
1// nx_pose_student_distill.nx -- BUILD-2 REAL DISTILLATION: our OWN sovereign student pose net learns to reproduce
2// the VALIDATED ViTPose TEACHER's pose on the REAL COCO image. Student = conv1(3->16,3x3,s2,p1)->ReLU->
3// conv2(16->48,3x3,s2,p1)->ReLU->conv3(48->17,1x1) -> 17 heatmaps [17,64,48] (4x downsample of the 256x192 input,
4// matching the teacher's output res). Input = the real preprocessed image (vitpose_input.bin). Targets = gaussian
5// bumps at the teacher's 17 keypoints (from the validated run). Trains via the gated loop (conv fwd/bwd + relu-bwd +
6// MSE + SGD), graded by nx_pck vs the teacher. Shows PCK 0 -> high = the student reproduces the teacher on a real
7// image = build-2 works at real scale (one-image overfit; generalization = train on many images = the arc).
8// expect_exit: 0
9import "nx_syscalls.nx"
10import "nx_f32.nx"
11import "nx_f32_cvt.nx"
12import "nx_f32_div.nx"
13import "nx_f32_exp.nx"
14import "nx_f32_conv2d.nx"
15import "nx_f32_conv2d_backward.nx"
16import "nx_f32_conv2d_fast.nx"
17import "nx_f32_conv2d_backward_fast.nx"
18import "nx_pose_cnn.nx"
19import "nx_f32_train_ops.nx"
20import "nx_pck.nx"
21const K_MAGIC_3072: i64 = 3072
22const K_MAGIC_12345: i64 = 12345
23const K_MAGIC_1103515245: i64 = 1103515245
24const K_MAGIC_2001: i64 = 2001
25const K_MAGIC_3600: i64 = 3600
26const K_MAGIC_8500: i64 = 8500
27const K_MAGIC_4900: i64 = 4900
28const K_MAGIC_300000: i64 = 300000
29
30func w(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 }
31func wn(v: i64) -> i64 { var m: i64=v; if m<0{w("-" as *u8);m=0-m} let t:*u8=sys_mmap(24); var k:i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i:i64=0; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1} sys_write(1,o,k); return 0 }
32func mk(n: i64, d: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(n), nx_i32_to_f32(d)) }
33// leaky ReLU (slope 0.01 for x<0) -- prevents the dying-ReLU collapse of the deep unnormalized net.
34func leaky_relu(buf: *i64, n: i64) -> i64 { let a: i64=mk(1,100); var i: i64=0; while i<n { if nx_f32_gt(buf[i],0)==0 { buf[i]=nx_f32_mul(a, buf[i]) } i=i+1 } return 0 }
35func leaky_relu_bwd(post: *i64, dpost: *i64, dpre: *i64, n: i64) -> i64 { let a: i64=mk(1,100); var i: i64=0; while i<n { if nx_f32_gt(post[i],0)==1 { dpre[i]=dpost[i] } else { dpre[i]=nx_f32_mul(a, dpost[i]) } i=i+1 } return 0 }
36// gradient clip to [-cap,cap] per element -- bounds weight growth, prevents the exploding-gradient collapse. Catches
37// Inf (nx_f32_gt(Inf,cap)==1); a NaN would slip through, so clipping from step 0 (while grads are still finite) is key.
38func f32_clip(buf: *i64, n: i64, cap: i64) -> i64 { let nc: i64=nx_f32_neg(cap); var i: i64=0; while i<n { if nx_f32_gt(buf[i],cap)==1 { buf[i]=cap } if nx_f32_lt(buf[i],nc)==1 { buf[i]=nc } i=i+1 } return 0 }
39
40// f32-aware argmax of heatmap j ([17,64,48], H=64 W=48) -> (ox,oy) in x0..47 y0..63
41func hm_argmax(heat: *i64, j: i64, ox: *i64, oy: *i64) -> i64 {
42 let hm: i64 = j*64*48
43 var best: i64 = heat[hm]; var bx: i64=0; var by: i64=0
44 var y: i64=0
45 while y<64 { var x: i64=0; while x<48 { let v: i64=heat[hm + y*48 + x]; if nx_f32_gt(v,best)==1 { best=v; bx=x; by=y } x=x+1 } y=y+1 }
46 ox[0]=bx; oy[0]=by; return 0
47}
48
49func main() -> i64 {
50 // ---- load the real preprocessed image [3,256,192] ----
51 let input: *i64=sys_mmap(8*3*256*192) as *i64
52 let iN: i64=3*256*192
53 let ifd: i64=sys_openat_rd("/home/elderwesto/vitpose_input.bin" as *u8)
54 if ifd<0 { w("no vitpose_input.bin\n" as *u8); sys_exit(1); return 1 }
55 let ib: *u8=sys_mmap(iN*4+16); var ig: i64=0
56 while ig<iN*4 { let n: i64=sys_read(ifd,((ib as i64)+ig) as *u8, iN*4-ig); if n<=0{ig=iN*4}else{ig=ig+n} }
57 sys_close(ifd)
58 var i: i64=0; while i<iN { input[i]=(ib[i*4]&0xff)|((ib[i*4+1]&0xff)<<8)|((ib[i*4+2]&0xff)<<16)|((ib[i*4+3]&0xff)<<24); i=i+1 }
59 w("loaded real image input\n" as *u8)
60
61 // ---- teacher keypoints (from the validated ViTPose run on this image) ----
62 let tkx: *i64=sys_mmap(8*17) as *i64; let tky: *i64=sys_mmap(8*17) as *i64
63 tkx[0]=26;tky[0]=11; tkx[1]=27;tky[1]=10; tkx[2]=26;tky[2]=10; tkx[3]=28;tky[3]=11; tkx[4]=25;tky[4]=11
64 tkx[5]=29;tky[5]=16; tkx[6]=26;tky[6]=17; tkx[7]=31;tky[7]=21; tkx[8]=25;tky[8]=23; tkx[9]=32;tky[9]=24
65 tkx[10]=22;tky[10]=26; tkx[11]=31;tky[11]=30; tkx[12]=29;tky[12]=31; tkx[13]=32;tky[13]=42; tkx[14]=27;tky[14]=41
66 tkx[15]=34;tky[15]=53; tkx[16]=29;tky[16]=51
67 // gaussian targets [17,64,48], sigma^2=4 -> exp(-d2/8)
68 let target: *i64=sys_mmap(8*17*64*48) as *i64
69 let ne: i64=mk(0-1,50) // sigma^2=25 -> WIDE gaussian (sharp targets -> background-dominated MSE -> collapse)
70 var j: i64=0
71 while j<17 { var y: i64=0; while y<64 { var x: i64=0; while x<48 {
72 let dx: i64=x-tkx[j]; let dy: i64=y-tky[j]; let d2: i64=dx*dx+dy*dy
73 target[j*K_MAGIC_3072 + y*48 + x]=nx_f32_exp(nx_f32_mul(nx_i32_to_f32(d2), ne))
74 x=x+1 } y=y+1 } j=j+1 }
75
76 // ---- student weights ----
77 let W1: *i64=sys_mmap(8*16*3*3*3) as *i64; let b1: *i64=sys_mmap(8*16) as *i64
78 let W2: *i64=sys_mmap(8*48*16*3*3) as *i64; let b2: *i64=sys_mmap(8*48) as *i64
79 let W3: *i64=sys_mmap(8*17*48) as *i64; let b3: *i64=sys_mmap(8*17) as *i64
80 var seed: i64=K_MAGIC_12345 // LCG random init at He scale (sqrt(2/fan_in)) -- breaks filter symmetry + avoids vanishing
81 i=0; while i<16*3*3*3 { seed=(seed*K_MAGIC_1103515245+K_MAGIC_12345)&0x7fffffff; W1[i]=mk((seed%K_MAGIC_2001)-1000, K_MAGIC_3600); i=i+1 } // ~+/-0.28 (fan27)
82 i=0; while i<48*16*3*3 { seed=(seed*K_MAGIC_1103515245+K_MAGIC_12345)&0x7fffffff; W2[i]=mk((seed%K_MAGIC_2001)-1000, K_MAGIC_8500); i=i+1 } // ~+/-0.12 (fan144)
83 i=0; while i<17*48 { seed=(seed*K_MAGIC_1103515245+K_MAGIC_12345)&0x7fffffff; W3[i]=mk((seed%K_MAGIC_2001)-1000, K_MAGIC_4900); i=i+1 } // ~+/-0.20 (fan48)
84 i=0; while i<16 { b1[i]=0; i=i+1 } i=0; while i<48 { b2[i]=0; i=i+1 } i=0; while i<17 { b3[i]=0; i=i+1 }
85
86 // buffers
87 let c1: *i64=sys_mmap(8*16*128*96) as *i64
88 let c2: *i64=sys_mmap(8*48*64*48) as *i64
89 let heat: *i64=sys_mmap(8*17*64*48) as *i64
90 let dHeat: *i64=sys_mmap(8*17*64*48) as *i64
91 let dA2: *i64=sys_mmap(8*48*64*48) as *i64; let dC2: *i64=sys_mmap(8*48*64*48) as *i64
92 let dA1: *i64=sys_mmap(8*16*128*96) as *i64; let dC1: *i64=sys_mmap(8*16*128*96) as *i64
93 let dW1: *i64=sys_mmap(8*16*3*3*3) as *i64; let db1: *i64=sys_mmap(8*16) as *i64
94 let dW2: *i64=sys_mmap(8*48*16*3*3) as *i64; let db2: *i64=sys_mmap(8*48) as *i64
95 let dW3: *i64=sys_mmap(8*17*48) as *i64; let db3: *i64=sys_mmap(8*17) as *i64
96 let dIn: *i64=sys_mmap(8*3*256*192) as *i64
97 let ox: *i64=sys_mmap(8) as *i64; let oy: *i64=sys_mmap(8) as *i64
98 let px: *i64=sys_mmap(8*17) as *i64; let py: *i64=sys_mmap(8*17) as *i64
99 let lr: i64=nx_q14_to_f32(160) // ~0.0098 (grad-clipped, so a working lr is safe)
100
101 // initial
102 nx_f32_conv2d_fast(input,1,3,256,192, W1,16,3,3,2,1, b1, c1); leaky_relu(c1, 16*128*96)
103 nx_f32_conv2d_fast(c1,1,16,128,96, W2,48,3,3,2,1, b2, c2); leaky_relu(c2, 48*64*48)
104 nx_f32_conv2d_fast(c2,1,48,64,48, W3,17,1,1,1,0, b3, heat)
105 let initL: i64=f32_mse_loss(heat, target, 17*64*48)
106 j=0; while j<17 { hm_argmax(heat,j,ox,oy); px[j]=ox[0]; py[j]=oy[0]; j=j+1 }
107 let initPCK: i64=pck_permille(px,py,tkx,tky,17,9) // within 3px
108 w("initial PCK@3px=" as *u8); wn(initPCK); w("\n" as *u8)
109
110 // train
111 let t0: i64=sys_now_us()
112 var step: i64=0
113 while step<300 {
114 nx_f32_conv2d_fast(input,1,3,256,192, W1,16,3,3,2,1, b1, c1); leaky_relu(c1, 16*128*96)
115 nx_f32_conv2d_fast(c1,1,16,128,96, W2,48,3,3,2,1, b2, c2); leaky_relu(c2, 48*64*48)
116 nx_f32_conv2d_fast(c2,1,48,64,48, W3,17,1,1,1,0, b3, heat)
117 f32_mse_grad(heat, target, dHeat, 17*64*48)
118 nx_f32_conv2d_backward_fast(c2,1,48,64,48, W3,17,1,1,1,0, dHeat, dA2, dW3, db3)
119 leaky_relu_bwd(c2, dA2, dC2, 48*64*48)
120 nx_f32_conv2d_backward_fast(c1,1,16,128,96, W2,48,3,3,2,1, dC2, dA1, dW2, db2)
121 leaky_relu_bwd(c1, dA1, dC1, 16*128*96)
122 nx_f32_conv2d_backward_fast(input,1,3,256,192, W1,16,3,3,2,1, dC1, dIn, dW1, db1)
123 let cap: i64=nx_i32_to_f32(1) // clip grads to +/-1 before SGD
124 f32_clip(dW1,16*3*3*3,cap); f32_clip(db1,16,cap); f32_clip(dW2,48*16*3*3,cap); f32_clip(db2,48,cap); f32_clip(dW3,17*48,cap); f32_clip(db3,17,cap)
125 f32_sgd_step(W1,dW1,lr,16*3*3*3); f32_sgd_step(b1,db1,lr,16)
126 f32_sgd_step(W2,dW2,lr,48*16*3*3); f32_sgd_step(b2,db2,lr,48)
127 f32_sgd_step(W3,dW3,lr,17*48); f32_sgd_step(b3,db3,lr,17)
128 if step % 30 == 0 {
129 w(" step " as *u8); wn(step)
130 j=0; while j<17 { hm_argmax(heat,j,ox,oy); px[j]=ox[0]; py[j]=oy[0]; j=j+1 }
131 w(" PCK@3=" as *u8); wn(pck_permille(px,py,tkx,tky,17,9)); w(" @5=" as *u8); wn(pck_permille(px,py,tkx,tky,17,25)); w("\n" as *u8)
132 }
133 step=step+1
134 }
135 let t1: i64=sys_now_us()
136 w(" TIMING: 300 steps in " as *u8); wn((t1-t0)/1000); w(" ms = " as *u8); wn((t1-t0)/K_MAGIC_300000); w(" ms/step\n" as *u8)
137
138 // final
139 nx_f32_conv2d_fast(input,1,3,256,192, W1,16,3,3,2,1, b1, c1); leaky_relu(c1, 16*128*96)
140 nx_f32_conv2d_fast(c1,1,16,128,96, W2,48,3,3,2,1, b2, c2); leaky_relu(c2, 48*64*48)
141 nx_f32_conv2d_fast(c2,1,48,64,48, W3,17,1,1,1,0, b3, heat)
142 let finalL: i64=f32_mse_loss(heat, target, 17*64*48)
143 j=0; while j<17 { hm_argmax(heat,j,ox,oy); px[j]=ox[0]; py[j]=oy[0]; j=j+1 }
144 let pck3: i64=pck_permille(px,py,tkx,tky,17,9)
145 let pck5: i64=pck_permille(px,py,tkx,tky,17,25)
146 w("final PCK@3px=" as *u8); wn(pck3); w(" PCK@5px=" as *u8); wn(pck5); w("\n" as *u8)
147 w(" student kpts: " as *u8); j=0; while j<17 { w("(" as *u8); wn(px[j]); w("," as *u8); wn(py[j]); w(")" as *u8); j=j+1 } w("\n" as *u8)
148
149 if initPCK < 300 { if pck5 >= 588 { // student learns from ~0 to matching >=10/17 within 5px
150 w("STUDENT-DISTILL OK: our sovereign student net reproduces the ViTPose teacher's pose on a real image (build-2 works)\n" as *u8)
151 sys_exit(0); return 0
152 } }
153 w("STUDENT-DISTILL partial: PCK@5=" as *u8); wn(pck5); w(" (loss initL vs finalL: see values)\n" as *u8)
154 sys_exit(1)
155 return 1
156}