nx_f32_bn_fold.nx source
↩ module page · 29 lines · 1565 B
1// nx_f32_bn_fold.nx -- fold inference BatchNorm into the PRECEDING conv's weights + bias (offline), so a ResNet/
2// HRNet-class pose net (which uses BN, not GroupNorm) loads straight into the existing dense/grouped f32 conv with
3// NO runtime BN kernel. For a conv output x = W*in + b followed by BN y = gamma*(x-mean)/sqrt(var+eps) + beta, the
4// fold is exact per output channel: s = gamma/sqrt(var+eps); W' = W*s; b' = (b-mean)*s + beta. Then foldedconv(in)
5// == BN(conv(in)) identically. Composes nx_f32_mul/add/sub/sqrt/div. This is the standard "BN-fold" every inference
6// runtime does; it means the porter never needs a live BatchNorm op. license_tier: ORIGINAL
7import "nx_syscalls.nx"
8import "nx_f32.nx"
9import "nx_f32_div.nx"
10
11// W [C_out, wpc] (wpc = C_in*KH*KW weights per output channel); b [C_out] or null. gamma/beta/mean/var [C_out].
12// eps is an f32 scalar. Writes W_out [C_out,wpc] and b_out [C_out].
13func nx_f32_bn_fold(W: *i64, b: *i64, C_out: i64, wpc: i64, gamma: *i64, beta: *i64, mean: *i64, vari: *i64, eps: i64, W_out: *i64, b_out: *i64) -> i64 {
14 var co: i64 = 0
15 while co < C_out {
16 let denom: i64 = nx_f32_sqrt(nx_f32_add(vari[co], eps))
17 let s: i64 = nx_f32_div(gamma[co], denom)
18 var bc: i64 = 0
19 if (b as i64) != 0 { bc = b[co] }
20 b_out[co] = nx_f32_add(nx_f32_mul(nx_f32_sub(bc, mean[co]), s), beta[co])
21 var k: i64 = 0
22 while k < wpc {
23 W_out[co * wpc + k] = nx_f32_mul(W[co * wpc + k], s)
24 k = k + 1
25 }
26 co = co + 1
27 }
28 return 0
29}