← Commits · b2b061d0
b2b061d00b14c1dd19744e582dcb7d386fee00da
diff --git a/test.c b/test.c
index e4c8bad..d97187d 100644
--- a/test.c
+++ b/test.c
@@ -269,6 +269,57 @@ static void test_layernorm_backward(ut_dev dev) {
ut_layernorm_free(&ln);
}
+static void test_batchnorm2d_forward(ut_dev dev) {
+ // x:[2,1,1,2] flat=[1,2,3,4], M=4, mu=2.5, var=1.25, rstd≈0.894424
+ ut_batchnorm2d bn = ut_batchnorm2d_alloc(1, dev);
+ ut_tensor* x = ut_from_data(4, (int[]){2, 1, 1, 2}, (float[]){1.f, 2.f, 3.f, 4.f}, dev);
+ ut_tensor* out = ut_batchnorm2d_forward(&bn, x, true, NULL);
+ ut_sync_cpu(out);
+ assert_data(out, ((float[]){-1.341635f, -0.447212f, 0.447212f, 1.341635f}), 1e-4f);
+
+ // running stats after one training step (momentum=0.1 default, unbiased var)
+ ut_sync_cpu(bn.running_mean);
+ ut_sync_cpu(bn.running_var);
+ assert_eq(bn.running_mean->data[0], 0.25f, 1e-4f);
+ assert_eq(bn.running_var->data[0], 1.066667f, 1e-4f);
+
+ // eval mode on a different input uses the running stats, not batch stats
+ ut_tensor* x2 = ut_from_data(4, (int[]){2, 1, 1, 2}, (float[]){10.f, 20.f, 30.f, 40.f}, dev);
+ ut_tensor* out2 = ut_batchnorm2d_forward(&bn, x2, false, NULL);
+ ut_sync_cpu(out2);
+ assert_data(out2, ((float[]){9.440353f, 19.122766f, 28.805179f, 38.487592f}), 1e-3f);
+
+ ut_free_all(x, out, x2, out2);
+ ut_batchnorm2d_free(&bn);
+}
+
+static void test_batchnorm2d_backward(ut_dev dev) {
+ // same x as the forward test; grad_out is one-hot [1,0,0,0]
+ ut_batchnorm2d bn = ut_batchnorm2d_alloc(1, dev);
+ ut_tensor* x = ut_from_data(4, (int[]){2, 1, 1, 2}, (float[]){1.f, 2.f, 3.f, 4.f}, dev);
+ ut_batchnorm2d_cache c;
+ ut_tensor* out = ut_batchnorm2d_forward(&bn, x, true, &c);
+ ut_free(out);
+
+ ut_tensor* go = ut_from_data(4, (int[]){2, 1, 1, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
+ ut_tensor* dW = ut_alloc(1, (int[]){1}, dev);
+ ut_tensor* db = ut_alloc(1, (int[]){1}, dev);
+ memset(dW->data, 0, sizeof(float));
+ memset(db->data, 0, sizeof(float));
+
+ ut_tensor* dx = ut_batchnorm2d_backward(&bn, &c, go, dW, db);
+ ut_sync_cpu(dW);
+ ut_sync_cpu(db);
+ ut_sync_cpu(dx);
+ assert_eq(dW->data[0], -1.341635f, 1e-4f);
+ assert_eq(db->data[0], 1.f, 1e-4f);
+ assert_data(dx, ((float[]){0.268330f, -0.357768f, -0.089443f, 0.178882f}), 1e-4f);
+
+ ut_free_all(dx, dW, db, go, x);
+ ut_batchnorm2d_cache_free(&c);
+ ut_batchnorm2d_free(&bn);
+}
+
static void test_im2col(ut_dev dev) {
// x[1,1,1,4]=[1,2,3,4], kh=1,kw=3,s=1,p=0 → col[3,2]
ut_tensor* x = ut_from_data(4, (int[]){1, 1, 1, 4}, (float[]){1.f, 2.f, 3.f, 4.f}, dev);
@@ -342,16 +393,15 @@ static void test_conv2d_forward(ut_dev dev) {
// oc0 W=[1,0;0,-1] bias=100, oc1 W=[0,1;1,0] bias=1000
ut_conv2d l = ut_conv2d_alloc(1, 2, 2, 2, 1, 0, true, dev);
ut_free(l.weight);
- l.weight = ut_from_data(4, (int[]){2, 1, 2, 2},
- (float[]){1.f, 0.f, 0.f, -1.f, 0.f, 1.f, 1.f, 0.f}, dev);
+ l.weight =
+ ut_from_data(4, (int[]){2, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, -1.f, 0.f, 1.f, 1.f, 0.f}, dev);
ut_free(l.bias);
l.bias = ut_from_data(1, (int[]){2}, (float[]){100.f, 1000.f}, dev);
- ut_tensor* x = ut_from_data(
- 4, (int[]){2, 1, 3, 3},
- (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 9.f, 8.f, 7.f, 6.f, 5.f, 4.f, 3.f, 2.f,
- 1.f},
- dev);
+ ut_tensor* x = ut_from_data(4, (int[]){2, 1, 3, 3},
+ (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f, 9.f, 8.f, 7.f,
+ 6.f, 5.f, 4.f, 3.f, 2.f, 1.f},
+ dev);
ut_tensor* out = ut_conv2d_forward(&l, x, NULL);
assert(out->dev == dev); // bias-add must not silently move a CPU model onto Metal
ut_sync_cpu(out);
@@ -381,8 +431,7 @@ static void test_conv2d_backward(ut_dev dev) {
ut_tensor* out = ut_conv2d_forward(&l, x, &c);
ut_free(out);
- ut_tensor* go =
- ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
+ ut_tensor* go = ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
ut_tensor* dW = ut_alloc(4, (int[]){1, 1, 2, 2}, dev);
memset(dW->data, 0, 16);
@@ -492,6 +541,11 @@ int main() {
test_layernorm_backward(UT_CPU);
test_layernorm_backward(UT_METAL);
+ test_batchnorm2d_forward(UT_CPU);
+ test_batchnorm2d_forward(UT_METAL);
+ test_batchnorm2d_backward(UT_CPU);
+ test_batchnorm2d_backward(UT_METAL);
+
test_im2col(UT_CPU);
test_im2col(UT_METAL);
test_col2im(UT_CPU);
diff --git a/utensil.h b/utensil.h
index 73ddd71..e61b209 100644
--- a/utensil.h
+++ b/utensil.h
@@ -69,6 +69,20 @@ typedef struct ut_conv2d_cache {
ut_tensor* col;
} ut_conv2d_cache;
+typedef struct ut_batchnorm2d {
+ ut_tensor* weight; // [C] gain
+ ut_tensor* bias; // [C] shift
+ ut_tensor* running_mean; // [C] EMA of batch mean, used at inference
+ ut_tensor* running_var; // [C] EMA of batch (unbiased) var, used at inference
+ int c; // channel count
+ float eps, momentum; // momentum: EMA weight given to each new batch stat
+} ut_batchnorm2d;
+
+typedef struct ut_batchnorm2d_cache {
+ ut_tensor* xnorm; // normalised x before affine
+ ut_tensor* rstd; // per-channel reciprocal std-dev
+} ut_batchnorm2d_cache;
+
typedef struct ut_sgd {
ut_tensor** params; // pointers to model parameters (not owned)
ut_tensor** grads; // gradient accumulators (owned)
@@ -1472,6 +1486,137 @@ void ut_conv2d_free(ut_conv2d* l) {
l->weight = l->bias = NULL;
}
+// =========================================================
+// BatchNorm2d
+// =========================================================
+ut_batchnorm2d ut_batchnorm2d_alloc(int c, ut_dev dev) {
+ ut_batchnorm2d l = {.c = c, .eps = 1e-5f, .momentum = 0.1f};
+ l.weight = ut_alloc(1, (int[]){c}, UT_CPU);
+ l.bias = ut_alloc(1, (int[]){c}, UT_CPU);
+ l.running_mean = ut_alloc(1, (int[]){c}, UT_CPU);
+ l.running_var = ut_alloc(1, (int[]){c}, UT_CPU);
+ for (int i = 0; i < c; i++) l.weight->data[i] = l.running_var->data[i] = 1.f;
+ if (dev == UT_METAL) {
+ ut_to_device(l.weight, UT_METAL);
+ ut_to_device(l.bias, UT_METAL);
+ ut_to_device(l.running_mean, UT_METAL);
+ ut_to_device(l.running_var, UT_METAL);
+ }
+ return l;
+}
+
+// x: [N,C,H,W] — normalises each channel over N,H,W. training=true uses batch
+// stats and updates running_mean/running_var (EMA); training=false normalises
+// with the running stats instead (eval/inference), matching nn.BatchNorm2d.
+ut_tensor* ut_batchnorm2d_forward(ut_batchnorm2d* l, ut_tensor* x, bool training,
+ ut_batchnorm2d_cache* cache) {
+ int N = x->shape.shape[0], C = l->c, H = x->shape.shape[2], W = x->shape.shape[3];
+ int M = N * H * W;
+ ut_sync_cpu(x);
+ ut_sync_cpu(l->weight);
+ ut_sync_cpu(l->bias);
+ ut_sync_cpu(l->running_mean);
+ ut_sync_cpu(l->running_var);
+
+ ut_tensor* out = ut_alloc(4, x->shape.shape, x->dev);
+ ut_tensor* xnorm_t = cache ? ut_alloc(4, x->shape.shape, UT_CPU) : NULL;
+ ut_tensor* rstd_t = cache ? ut_alloc(1, (int[]){C}, UT_CPU) : NULL;
+
+ for (int c = 0; c < C; c++) {
+ float mu, rs;
+ if (training) {
+ float sum = 0;
+ for (int n = 0; n < N; n++)
+ for (int h = 0; h < H; h++)
+ for (int w = 0; w < W; w++) sum += x->data[((n * C + c) * H + h) * W + w];
+ mu = sum / (float)M;
+ float var = 0;
+ for (int n = 0; n < N; n++)
+ for (int h = 0; h < H; h++)
+ for (int w = 0; w < W; w++) {
+ float diff = x->data[((n * C + c) * H + h) * W + w] - mu;
+ var += diff * diff;
+ }
+ var /= (float)M;
+ rs = 1.f / sqrtf(var + l->eps);
+ float var_unbiased = M > 1 ? var * (float)M / (float)(M - 1) : var;
+ l->running_mean->data[c] = (1.f - l->momentum) * l->running_mean->data[c] + l->momentum * mu;
+ l->running_var->data[c] =
+ (1.f - l->momentum) * l->running_var->data[c] + l->momentum * var_unbiased;
+ } else {
+ mu = l->running_mean->data[c];
+ rs = 1.f / sqrtf(l->running_var->data[c] + l->eps);
+ }
+ for (int n = 0; n < N; n++)
+ for (int h = 0; h < H; h++)
+ for (int w = 0; w < W; w++) {
+ int idx = ((n * C + c) * H + h) * W + w;
+ float xn = (x->data[idx] - mu) * rs;
+ out->data[idx] = l->weight->data[c] * xn + l->bias->data[c];
+ if (cache) xnorm_t->data[idx] = xn;
+ }
+ if (cache) rstd_t->data[c] = rs;
+ }
+ out->dirty_gpu = true;
+ if (training) {
+ l->running_mean->dirty_gpu = true;
+ l->running_var->dirty_gpu = true;
+ }
+ if (cache) {
+ cache->xnorm = xnorm_t;
+ cache->rstd = rstd_t;
+ }
+ return out;
+}
+
+ut_tensor* ut_batchnorm2d_backward(ut_batchnorm2d* l, ut_batchnorm2d_cache* cache,
+ ut_tensor* grad_out, ut_tensor* dW, ut_tensor* db) {
+ int N = grad_out->shape.shape[0], C = l->c, H = grad_out->shape.shape[2],
+ W = grad_out->shape.shape[3];
+ int M = N * H * W;
+ ut_sync_cpu(grad_out);
+ ut_sync_cpu(cache->xnorm);
+ ut_sync_cpu(dW);
+ ut_sync_cpu(db);
+ ut_sync_cpu(l->weight);
+
+ ut_tensor* dx = ut_alloc(4, grad_out->shape.shape, grad_out->dev);
+ for (int c = 0; c < C; c++) {
+ float sum_go = 0, sum_go_xn = 0;
+ for (int n = 0; n < N; n++)
+ for (int h = 0; h < H; h++)
+ for (int w = 0; w < W; w++) {
+ int idx = ((n * C + c) * H + h) * W + w;
+ sum_go += grad_out->data[idx];
+ sum_go_xn += grad_out->data[idx] * cache->xnorm->data[idx];
+ }
+ dW->data[c] += sum_go_xn;
+ db->data[c] += sum_go;
+
+ float rs = cache->rstd->data[c], wgt = l->weight->data[c];
+ for (int n = 0; n < N; n++)
+ for (int h = 0; h < H; h++)
+ for (int w = 0; w < W; w++) {
+ int idx = ((n * C + c) * H + h) * W + w;
+ float go = grad_out->data[idx], xn = cache->xnorm->data[idx];
+ dx->data[idx] = rs * wgt * (go - (sum_go + xn * sum_go_xn) / (float)M);
+ }
+ }
+ dW->dirty_gpu = true;
+ db->dirty_gpu = true;
+ dx->dirty_gpu = true;
+ return dx;
+}
+
+void ut_batchnorm2d_cache_free(ut_batchnorm2d_cache* c) {
+ ut_free_all(c->xnorm, c->rstd);
+ c->xnorm = c->rstd = NULL;
+}
+void ut_batchnorm2d_free(ut_batchnorm2d* l) {
+ ut_free_all(l->weight, l->bias, l->running_mean, l->running_var);
+ l->weight = l->bias = l->running_mean = l->running_var = NULL;
+}
+
// =========================================================
// Softmax
// =========================================================