← Commits · 1eb11f48
1eb11f48d1005ae02a80d09567d6f7ce8d0267a1
diff --git a/.gitignore b/.gitignore
index dfdb3e0..672c3aa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,3 +6,6 @@ test
mnist/
examples/mnist_mlp
examples/mnist_cnn
+examples/mobilenetv1_cifar10
+examples/resnet8_cifar10
+
diff --git a/Makefile b/Makefile
index c32692f..c82f5e1 100644
--- a/Makefile
+++ b/Makefile
@@ -4,7 +4,7 @@ LDFLAGS ?= -lm
METAL := $(shell test -d /System/Library/Frameworks/Metal.framework && echo 1 || echo 0)
ifeq ($(METAL),1)
CFLAGS += -DACCELERATE_NEW_LAPACK
- LDFLAGS += -framework Metal -framework MetalPerformanceShaders -framework Foundation -framework Accelerate
+ LDFLAGS += -framework Metal -framework MetalPerformanceShaders -framework MetalPerformanceShadersGraph -framework Foundation -framework Accelerate
endif
test:
diff --git a/examples/cifar10_loader.h b/examples/cifar10_loader.h
new file mode 100644
index 0000000..60b1b1e
--- /dev/null
+++ b/examples/cifar10_loader.h
@@ -0,0 +1,71 @@
+#ifndef CIFAR10_LOADER_H
+#define CIFAR10_LOADER_H
+
+// Minimal loader for the CIFAR-10 "binary version" (cifar-10-binary.tar.gz):
+// each record is 1 label byte + 3072 pixel bytes (1024 R, 1024 G, 1024 B; each
+// plane row-major 32x32). Include after utensil.h (cifar10_argmax takes a
+// ut_tensor*). Normalises with the standard CIFAR-10 per-channel mean/std.
+
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#define CIFAR10_PIX (3 * 32 * 32)
+#define CIFAR10_REC (1 + CIFAR10_PIX)
+
+typedef struct {
+ int n;
+ float* imgs; // [n,3,32,32] NCHW, normalised
+ uint8_t* labels; // [n]
+} cifar10_t;
+
+static const float CIFAR10_MEAN[3] = {0.4914f, 0.4822f, 0.4465f};
+static const float CIFAR10_STD[3] = {0.2470f, 0.2435f, 0.2616f};
+
+// paths: nfiles binary batch files (data_batch_1..5.bin, or test_batch.bin),
+// each holding exactly 10000 records
+static cifar10_t cifar10_load(const char** paths, int nfiles) {
+ cifar10_t d = {.n = nfiles * 10000};
+ d.imgs = malloc((size_t)d.n * CIFAR10_PIX * sizeof(float));
+ d.labels = malloc((size_t)d.n);
+ uint8_t rec[CIFAR10_REC];
+ int idx = 0;
+ for (int f = 0; f < nfiles; f++) {
+ FILE* fp = fopen(paths[f], "rb");
+ for (int i = 0; i < 10000; i++) {
+ fread(rec, 1, CIFAR10_REC, fp);
+ d.labels[idx] = rec[0];
+ for (int p = 0; p < CIFAR10_PIX; p++) {
+ int ch = p / (32 * 32);
+ d.imgs[idx * CIFAR10_PIX + p] = (rec[1 + p] / 255.f - CIFAR10_MEAN[ch]) / CIFAR10_STD[ch];
+ }
+ idx++;
+ }
+ fclose(fp);
+ }
+ return d;
+}
+
+static void cifar10_free(cifar10_t* d) {
+ free(d->imgs);
+ free(d->labels);
+}
+
+static void cifar10_shuffle(int* a, int n) {
+ for (int i = n - 1; i > 0; i--) {
+ int j = rand() % (i + 1);
+ int t = a[i];
+ a[i] = a[j];
+ a[j] = t;
+ }
+}
+
+static int cifar10_argmax(ut_tensor* t, int row) {
+ int C = t->shape.shape[1], best = 0;
+ float* r = t->data + row * C;
+ for (int j = 1; j < C; j++)
+ if (r[j] > r[best]) best = j;
+ return best;
+}
+
+#endif
diff --git a/examples/mobilenetv1_cifar10.c b/examples/mobilenetv1_cifar10.c
new file mode 100644
index 0000000..04d1559
--- /dev/null
+++ b/examples/mobilenetv1_cifar10.c
@@ -0,0 +1,305 @@
+#include "utensil.h"
+
+#include "cifar10_loader.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+
+// MobileNetV1 (Howard et al.), scaled down for CIFAR-10 (32x32, 10 classes)
+// instead of the ImageNet table (224x224, 1000 classes, channels up to 1024) --
+// that table doesn't fit a 32x32 input and would be wildly over-parameterized
+// for a 10-class problem. Keeps the architecture's defining idea (depthwise
+// separable conv: a per-channel spatial conv + a 1x1 pointwise conv, each
+// followed by BN+ReLU6, replacing one dense conv) and its downsampling
+// pattern (stride-2 in a depthwise conv, never in the pointwise conv), just
+// right-sized: stem -> 6 depthwise-separable blocks (channels 32-256, three
+// stride-2 steps: 32x32 -> 16x16 -> 8x8 -> 4x4) -> GAP -> linear head. No
+// data augmentation -- a comparison scaffold, not a from-scratch SOTA run.
+
+// gradient tensors are borrowed from the optimizer's own grads[]; wired once
+// in main() after ut_adam_alloc, then read directly during backward.
+typedef struct {
+ ut_tensor *dw_w, *dw_b, *dw_bn_w, *dw_bn_b, *pw_w, *pw_b, *pw_bn_w, *pw_bn_b;
+} block_grads_t;
+
+typedef struct {
+ ut_dwconv2d dw; // depthwise: one 3x3 filter per input channel
+ ut_batchnorm2d dw_bn;
+ ut_conv2d pw; // pointwise: 1x1 conv, mixes channels
+ ut_batchnorm2d pw_bn;
+ block_grads_t g;
+} block_t;
+
+typedef struct {
+ ut_dwconv2d_cache dw_c;
+ ut_batchnorm2d_cache dw_bn_c;
+ ut_conv2d_cache pw_c;
+ ut_batchnorm2d_cache pw_bn_c;
+ ut_tensor *dw_relu_in, *pw_relu_in; // kept for backward's two ReLU6es
+} block_cache_t;
+
+typedef struct {
+ ut_conv2d stem_conv;
+ ut_batchnorm2d stem_bn;
+ block_t blocks[6];
+ ut_linear fc;
+ ut_tensor *stem_cw, *stem_cb, *stem_bw, *stem_bb, *fc_w, *fc_b; // borrowed, like block_grads_t
+} mobilenet_t;
+
+typedef struct {
+ ut_conv2d_cache stem_c;
+ ut_batchnorm2d_cache stem_bn_c;
+ ut_tensor* stem_relu_in;
+ block_cache_t block_c[6];
+ ut_tensor* pooled; // kept for the fc layer's backward
+ int pool_h, pool_w;
+} mobilenet_cache_t;
+
+static block_t block_alloc(int in_c, int out_c, int stride, ut_dev dev) {
+ block_t b = {0};
+ b.dw = ut_dwconv2d_alloc(in_c, 3, 3, stride, 1, true, dev);
+ b.dw_bn = ut_batchnorm2d_alloc(in_c, dev);
+ b.pw = ut_conv2d_alloc(in_c, out_c, 1, 1, 1, 0, true, dev);
+ b.pw_bn = ut_batchnorm2d_alloc(out_c, dev);
+ return b;
+}
+
+static mobilenet_t mobilenet_alloc(ut_dev dev) {
+ mobilenet_t m = {0};
+ m.stem_conv = ut_conv2d_alloc(3, 32, 3, 3, 1, 1, true, dev);
+ m.stem_bn = ut_batchnorm2d_alloc(32, dev);
+ m.blocks[0] = block_alloc(32, 64, 1, dev); // 32x32
+ m.blocks[1] = block_alloc(64, 128, 2, dev); // 32x32 -> 16x16
+ m.blocks[2] = block_alloc(128, 128, 1, dev); // 16x16
+ m.blocks[3] = block_alloc(128, 256, 2, dev); // 16x16 -> 8x8
+ m.blocks[4] = block_alloc(256, 256, 1, dev); // 8x8
+ m.blocks[5] = block_alloc(256, 256, 2, dev); // 8x8 -> 4x4
+ m.fc = ut_linear_alloc(256, 10, true, dev);
+ return m;
+}
+
+// appends this block's param tensors to out[] (for ut_adam_alloc) and returns
+// the count -- always 8, no conditional (unlike a residual block's shortcut)
+static int block_params(block_t* b, ut_tensor** out) {
+ int n = 0;
+ out[n++] = b->dw.weight, out[n++] = b->dw.bias;
+ out[n++] = b->dw_bn.weight, out[n++] = b->dw_bn.bias;
+ out[n++] = b->pw.weight, out[n++] = b->pw.bias;
+ out[n++] = b->pw_bn.weight, out[n++] = b->pw_bn.bias;
+ return n;
+}
+
+// mirrors block_params' exact order to wire each gradient pointer to its param
+static void block_wire_grads(block_t* b, ut_tensor** g) {
+ int n = 0;
+ b->g.dw_w = g[n++], b->g.dw_b = g[n++];
+ b->g.dw_bn_w = g[n++], b->g.dw_bn_b = g[n++];
+ b->g.pw_w = g[n++], b->g.pw_b = g[n++];
+ b->g.pw_bn_w = g[n++], b->g.pw_bn_b = g[n++];
+}
+
+static ut_tensor* block_forward(block_t* b, ut_tensor* x, bool training, block_cache_t* c) {
+ ut_tensor* d1 = ut_dwconv2d_forward(&b->dw, x, c ? &c->dw_c : NULL);
+ ut_tensor* d2 = ut_batchnorm2d_forward(&b->dw_bn, d1, training, c ? &c->dw_bn_c : NULL);
+ ut_free(d1);
+ ut_tensor* d3 = ut_relu6(d2);
+ if (c) c->dw_relu_in = d2; else ut_free(d2);
+
+ ut_tensor* p1 = ut_conv2d_forward(&b->pw, d3, c ? &c->pw_c : NULL);
+ ut_free(d3);
+ ut_tensor* p2 = ut_batchnorm2d_forward(&b->pw_bn, p1, training, c ? &c->pw_bn_c : NULL);
+ ut_free(p1);
+ ut_tensor* out = ut_relu6(p2);
+ if (c) c->pw_relu_in = p2; else ut_free(p2);
+ return out;
+}
+
+static ut_tensor* block_backward(block_t* b, block_cache_t* c, ut_tensor* dout) {
+ ut_tensor* dp2 = ut_relu6_backward(dout, c->pw_relu_in);
+ ut_tensor* dp1 = ut_batchnorm2d_backward(&b->pw_bn, &c->pw_bn_c, dp2, b->g.pw_bn_w, b->g.pw_bn_b);
+ ut_free(dp2);
+ ut_tensor* dd3 = ut_conv2d_backward(&b->pw, &c->pw_c, dp1, b->g.pw_w, b->g.pw_b);
+ ut_free(dp1);
+ ut_tensor* dd2 = ut_relu6_backward(dd3, c->dw_relu_in);
+ ut_free(dd3);
+ ut_tensor* dd1 = ut_batchnorm2d_backward(&b->dw_bn, &c->dw_bn_c, dd2, b->g.dw_bn_w, b->g.dw_bn_b);
+ ut_free(dd2);
+ ut_tensor* dx = ut_dwconv2d_backward(&b->dw, &c->dw_c, dd1, b->g.dw_w, b->g.dw_b);
+ ut_free(dd1);
+ return dx;
+}
+
+static void block_cache_free(block_cache_t* c) {
+ ut_dwconv2d_cache_free(&c->dw_c);
+ ut_batchnorm2d_cache_free(&c->dw_bn_c);
+ ut_conv2d_cache_free(&c->pw_c);
+ ut_batchnorm2d_cache_free(&c->pw_bn_c);
+ ut_free_all(c->dw_relu_in, c->pw_relu_in);
+}
+
+static void block_free(block_t* b) {
+ ut_dwconv2d_free(&b->dw);
+ ut_batchnorm2d_free(&b->dw_bn);
+ ut_conv2d_free(&b->pw);
+ ut_batchnorm2d_free(&b->pw_bn);
+}
+
+static ut_tensor* mobilenet_forward(mobilenet_t* m, ut_tensor* x, bool training,
+ mobilenet_cache_t* c) {
+ ut_tensor* s = ut_conv2d_forward(&m->stem_conv, x, c ? &c->stem_c : NULL);
+ ut_tensor* sbn = ut_batchnorm2d_forward(&m->stem_bn, s, training, c ? &c->stem_bn_c : NULL);
+ ut_free(s);
+ ut_tensor* h = ut_relu6(sbn);
+ if (c) c->stem_relu_in = sbn; else ut_free(sbn);
+
+ for (int i = 0; i < 6; i++) {
+ ut_tensor* next = block_forward(&m->blocks[i], h, training, c ? &c->block_c[i] : NULL);
+ ut_free(h);
+ h = next;
+ }
+
+ int ph = h->shape.shape[2], pw = h->shape.shape[3];
+ ut_tensor* pooled = ut_global_avgpool2d(h);
+ ut_free(h);
+ if (c) c->pooled = pooled, c->pool_h = ph, c->pool_w = pw;
+ ut_tensor* logits = ut_linear_forward(&m->fc, pooled);
+ if (!c) ut_free(pooled);
+ return logits;
+}
+
+static ut_tensor* mobilenet_backward(mobilenet_t* m, mobilenet_cache_t* c, ut_tensor* dlogits) {
+ ut_tensor* dpooled = ut_linear_backward(&m->fc, c->pooled, dlogits, m->fc_w, m->fc_b);
+ ut_tensor* dh = ut_global_avgpool2d_backward(dpooled, c->pool_h, c->pool_w);
+ ut_free(dpooled);
+
+ for (int i = 5; i >= 0; i--) {
+ ut_tensor* dprev = block_backward(&m->blocks[i], &c->block_c[i], dh);
+ ut_free(dh);
+ dh = dprev;
+ }
+
+ ut_tensor* dsbn = ut_relu6_backward(dh, c->stem_relu_in);
+ ut_free(dh);
+ ut_tensor* ds = ut_batchnorm2d_backward(&m->stem_bn, &c->stem_bn_c, dsbn, m->stem_bw, m->stem_bb);
+ ut_free(dsbn);
+ ut_tensor* dx = ut_conv2d_backward(&m->stem_conv, &c->stem_c, ds, m->stem_cw, m->stem_cb);
+ ut_free(ds);
+ return dx;
+}
+
+static void mobilenet_cache_free(mobilenet_cache_t* c) {
+ ut_conv2d_cache_free(&c->stem_c);
+ ut_batchnorm2d_cache_free(&c->stem_bn_c);
+ ut_free(c->stem_relu_in);
+ for (int i = 0; i < 6; i++) block_cache_free(&c->block_c[i]);
+ ut_free(c->pooled);
+}
+
+static void mobilenet_free(mobilenet_t* m) {
+ ut_conv2d_free(&m->stem_conv);
+ ut_batchnorm2d_free(&m->stem_bn);
+ for (int i = 0; i < 6; i++) block_free(&m->blocks[i]);
+ ut_linear_free(&m->fc);
+}
+
+int main(void) {
+ srand(42);
+
+ ut_dev dev = UT_METAL; // flip to UT_CPU to compare against a (CPU-only) PyTorch script
+
+ printf("Loading CIFAR-10…\n");
+ const char* train_files[5] = {"cifar/data_batch_1.bin", "cifar/data_batch_2.bin",
+ "cifar/data_batch_3.bin", "cifar/data_batch_4.bin",
+ "cifar/data_batch_5.bin"};
+ const char* test_files[1] = {"cifar/test_batch.bin"};
+ cifar10_t train = cifar10_load(train_files, 5);
+ cifar10_t test = cifar10_load(test_files, 1);
+ printf("train: %d test: %d\n\n", train.n, test.n);
+
+ mobilenet_t m = mobilenet_alloc(dev);
+
+ ut_tensor* params[64];
+ int np = 0;
+ params[np++] = m.stem_conv.weight, params[np++] = m.stem_conv.bias;
+ params[np++] = m.stem_bn.weight, params[np++] = m.stem_bn.bias;
+ int block_off[6];
+ for (int i = 0; i < 6; i++) {
+ block_off[i] = np;
+ np += block_params(&m.blocks[i], params + np);
+ }
+ int fc_off = np;
+ params[np++] = m.fc.weight, params[np++] = m.fc.bias;
+
+ ut_adam opt = ut_adam_alloc(params, np, 1e-3f, 0.9f, 0.999f, 1e-8f, 1e-4f);
+
+ // wire each layer's gradient pointers to the optimizer's own grads[], in the
+ // same order the params[] above was built
+ m.stem_cw = opt.grads[0], m.stem_cb = opt.grads[1];
+ m.stem_bw = opt.grads[2], m.stem_bb = opt.grads[3];
+ for (int i = 0; i < 6; i++) block_wire_grads(&m.blocks[i], opt.grads + block_off[i]);
+ m.fc_w = opt.grads[fc_off], m.fc_b = opt.grads[fc_off + 1];
+
+ int B = 64, epochs = 20, batches = train.n / B;
+ ut_tensor* grad_logits = ut_alloc(2, (int[]){B, 10}, dev);
+ int* idx = malloc((size_t)train.n * sizeof(int));
+ for (int i = 0; i < train.n; i++) idx[i] = i;
+
+ for (int ep = 0; ep < epochs; ep++) {
+ cifar10_shuffle(idx, train.n);
+ clock_t t0 = clock();
+ float loss_sum = 0;
+ int correct = 0;
+
+ for (int bi = 0; bi < batches; bi++) {
+ float bx[B * CIFAR10_PIX];
+ int bl[B];
+ for (int i = 0; i < B; i++) {
+ int ii = idx[bi * B + i];
+ memcpy(bx + i * CIFAR10_PIX, train.imgs + ii * CIFAR10_PIX, CIFAR10_PIX * sizeof(float));
+ bl[i] = train.labels[ii];
+ }
+ ut_tensor* x = ut_from_data(4, (int[]){B, 3, 32, 32}, bx, dev);
+
+ mobilenet_cache_t c;
+ ut_tensor* logits = mobilenet_forward(&m, x, true, &c);
+ float loss = ut_cross_entropy(logits, bl, grad_logits);
+ ut_sync_cpu(logits);
+ for (int i = 0; i < B; i++)
+ if (cifar10_argmax(logits, i) == bl[i]) correct++;
+
+ ut_tensor* dx = mobilenet_backward(&m, &c, grad_logits);
+ ut_adam_step(&opt, 5.0f);
+ loss_sum += loss;
+
+ mobilenet_cache_free(&c);
+ ut_free_all(x, logits, dx);
+ }
+
+ float secs = (float)(clock() - t0) / (float)CLOCKS_PER_SEC;
+
+ int test_ok = 0;
+ for (int i = 0; i < test.n; i += B) {
+ int nb = (i + B <= test.n) ? B : (test.n - i);
+ ut_tensor* tx = ut_from_data(4, (int[]){nb, 3, 32, 32}, test.imgs + i * CIFAR10_PIX, dev);
+ ut_tensor* tl = mobilenet_forward(&m, tx, false, NULL);
+ ut_sync_cpu(tl);
+ for (int j = 0; j < nb; j++)
+ if (cifar10_argmax(tl, j) == test.labels[i + j]) test_ok++;
+ ut_free_all(tx, tl);
+ }
+
+ printf("epoch %2d loss %7.4f train %5.1f%% test %5.1f%% %6.1fs\n", ep + 1,
+ loss_sum / (float)batches, 100.f * (float)correct / (float)(batches * B),
+ 100.f * (float)test_ok / (float)test.n, secs);
+ }
+
+ free(idx);
+ ut_free(grad_logits);
+ ut_adam_free(&opt);
+ mobilenet_free(&m);
+ cifar10_free(&train);
+ cifar10_free(&test);
+ return 0;
+}
diff --git a/examples/resnet8_cifar10.c b/examples/resnet8_cifar10.c
new file mode 100644
index 0000000..80297ab
--- /dev/null
+++ b/examples/resnet8_cifar10.c
@@ -0,0 +1,339 @@
+#include "utensil.h"
+
+#include "cifar10_loader.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <time.h>
+
+// Standard CIFAR-10 ResNet-8 (He et al.'s CIFAR ResNet family, n=1): a stem
+// conv, then 3 residual stages (16/32/64 channels, stride 1/2/2) of one basic
+// block each, GAP, and a linear head. Stride/channel-changing blocks use a 1x1
+// conv+BN projection shortcut. No data augmentation, so expect lower accuracy
+// than a fully-tuned ResNet-8 — this is a scaffold for comparing against the
+// same architecture in PyTorch, not a from-scratch SOTA run.
+
+// gradient tensors are borrowed from the optimizer's own grads[]; wired once
+// in main() after ut_adam_alloc, then read directly during backward.
+typedef struct {
+ ut_tensor *c1w, *c1b, *b1w, *b1b, *c2w, *c2b, *b2w, *b2b, *pw, *pb, *pbnw, *pbnb;
+} block_grads_t;
+
+typedef struct {
+ ut_conv2d conv1, conv2, proj;
+ ut_batchnorm2d bn1, bn2, projbn;
+ bool has_proj;
+ block_grads_t g;
+} block_t;
+
+typedef struct {
+ ut_conv2d_cache c1, c2, cp;
+ ut_batchnorm2d_cache b1, b2, bp;
+ ut_tensor *h1bn, *sum; // kept for backward's two ReLUs
+} block_cache_t;
+
+typedef struct {
+ ut_conv2d stem_conv;
+ ut_batchnorm2d stem_bn;
+ block_t blocks[3];
+ ut_linear fc;
+ ut_tensor *stem_cw, *stem_cb, *stem_bw, *stem_bb, *fc_w, *fc_b; // borrowed, like block_grads_t
+} resnet8_t;
+
+typedef struct {
+ ut_conv2d_cache stem_c;
+ ut_batchnorm2d_cache stem_bn_c;
+ ut_tensor* stem_relu_in;
+ block_cache_t block_c[3];
+ ut_tensor* pooled; // kept for the fc layer's backward
+ int pool_h, pool_w;
+} resnet8_cache_t;
+
+static block_t block_alloc(int in_c, int out_c, int stride, ut_dev dev) {
+ block_t b = {0};
+ b.conv1 = ut_conv2d_alloc(in_c, out_c, 3, 3, stride, 1, true, dev);
+ b.bn1 = ut_batchnorm2d_alloc(out_c, dev);
+ b.conv2 = ut_conv2d_alloc(out_c, out_c, 3, 3, 1, 1, true, dev);
+ b.bn2 = ut_batchnorm2d_alloc(out_c, dev);
+ b.has_proj = in_c != out_c || stride != 1;
+ if (b.has_proj) {
+ b.proj = ut_conv2d_alloc(in_c, out_c, 1, 1, stride, 0, true, dev);
+ b.projbn = ut_batchnorm2d_alloc(out_c, dev);
+ }
+ return b;
+}
+
+static resnet8_t resnet8_alloc(ut_dev dev) {
+ resnet8_t m = {0};
+ m.stem_conv = ut_conv2d_alloc(3, 16, 3, 3, 1, 1, true, dev);
+ m.stem_bn = ut_batchnorm2d_alloc(16, dev);
+ m.blocks[0] = block_alloc(16, 16, 1, dev);
+ m.blocks[1] = block_alloc(16, 32, 2, dev);
+ m.blocks[2] = block_alloc(32, 64, 2, dev);
+ m.fc = ut_linear_alloc(64, 10, true, dev);
+ return m;
+}
+
+// appends this block's param tensors to out[] (for ut_adam_alloc) and returns the count
+static int block_params(block_t* b, ut_tensor** out) {
+ int n = 0;
+ out[n++] = b->conv1.weight, out[n++] = b->conv1.bias;
+ out[n++] = b->bn1.weight, out[n++] = b->bn1.bias;
+ out[n++] = b->conv2.weight, out[n++] = b->conv2.bias;
+ out[n++] = b->bn2.weight, out[n++] = b->bn2.bias;
+ if (b->has_proj) {
+ out[n++] = b->proj.weight, out[n++] = b->proj.bias;
+ out[n++] = b->projbn.weight, out[n++] = b->projbn.bias;
+ }
+ return n;
+}
+
+// mirrors block_params' exact order to wire each gradient pointer to its param
+static void block_wire_grads(block_t* b, ut_tensor** g) {
+ int n = 0;
+ b->g.c1w = g[n++], b->g.c1b = g[n++];
+ b->g.b1w = g[n++], b->g.b1b = g[n++];
+ b->g.c2w = g[n++], b->g.c2b = g[n++];
+ b->g.b2w = g[n++], b->g.b2b = g[n++];
+ if (b->has_proj) {
+ b->g.pw = g[n++], b->g.pb = g[n++];
+ b->g.pbnw = g[n++], b->g.pbnb = g[n++];
+ }
+}
+
+static ut_tensor* block_forward(block_t* b, ut_tensor* x, bool training, block_cache_t* c) {
+ ut_tensor* h1 = ut_conv2d_forward(&b->conv1, x, c ? &c->c1 : NULL);
+ ut_tensor* h1bn = ut_batchnorm2d_forward(&b->bn1, h1, training, c ? &c->b1 : NULL);
+ ut_free(h1);
+ ut_tensor* h1r = ut_relu(h1bn);
+ if (c) c->h1bn = h1bn; else ut_free(h1bn);
+
+ ut_tensor* h2 = ut_conv2d_forward(&b->conv2, h1r, c ? &c->c2 : NULL);
+ ut_free(h1r);
+ ut_tensor* h2bn = ut_batchnorm2d_forward(&b->bn2, h2, training, c ? &c->b2 : NULL);
+ ut_free(h2);
+
+ ut_tensor* shortcut;
+ if (b->has_proj) {
+ ut_tensor* p = ut_conv2d_forward(&b->proj, x, c ? &c->cp : NULL);
+ shortcut = ut_batchnorm2d_forward(&b->projbn, p, training, c ? &c->bp : NULL);
+ ut_free(p);
+ } else {
+ shortcut = ut_retain(x);
+ }
+
+ ut_tensor* sum = ut_add(h2bn, shortcut);
+ ut_free_all(h2bn, shortcut);
+ ut_tensor* out = ut_relu(sum);
+ if (c) c->sum = sum; else ut_free(sum);
+ return out;
+}
+
+static ut_tensor* block_backward(block_t* b, block_cache_t* c, ut_tensor* dout) {
+ ut_tensor* dsum = ut_relu_backward(dout, c->sum);
+ // sum = h2bn + shortcut -> gradient passes unchanged to both branches
+ ut_tensor* dh2 = ut_batchnorm2d_backward(&b->bn2, &c->b2, dsum, b->g.b2w, b->g.b2b);
+ ut_tensor* dh1r = ut_conv2d_backward(&b->conv2, &c->c2, dh2, b->g.c2w, b->g.c2b);
+ ut_free(dh2);
+ ut_tensor* dh1bn = ut_relu_backward(dh1r, c->h1bn);
+ ut_free(dh1r);
+ ut_tensor* dh1 = ut_batchnorm2d_backward(&b->bn1, &c->b1, dh1bn, b->g.b1w, b->g.b1b);
+ ut_free(dh1bn);
+ ut_tensor* dx_main = ut_conv2d_backward(&b->conv1, &c->c1, dh1, b->g.c1w, b->g.c1b);
+ ut_free(dh1);
+
+ ut_tensor* dx;
+ if (b->has_proj) {
+ ut_tensor* dp = ut_batchnorm2d_backward(&b->projbn, &c->bp, dsum, b->g.pbnw, b->g.pbnb);
+ ut_tensor* dx_proj = ut_conv2d_backward(&b->proj, &c->cp, dp, b->g.pw, b->g.pb);
+ ut_free(dp);
+ dx = ut_add(dx_main, dx_proj);
+ ut_free_all(dx_main, dx_proj);
+ } else {
+ dx = ut_add(dx_main, dsum);
+ ut_free(dx_main);
+ }
+ ut_free(dsum);
+ return dx;
+}
+
+static void block_cache_free(block_t* b, block_cache_t* c) {
+ ut_conv2d_cache_free(&c->c1);
+ ut_conv2d_cache_free(&c->c2);
+ ut_batchnorm2d_cache_free(&c->b1);
+ ut_batchnorm2d_cache_free(&c->b2);
+ if (b->has_proj) {
+ ut_conv2d_cache_free(&c->cp);
+ ut_batchnorm2d_cache_free(&c->bp);
+ }
+ ut_free_all(c->h1bn, c->sum);
+}
+
+static void block_free(block_t* b) {
+ ut_conv2d_free(&b->conv1);
+ ut_conv2d_free(&b->conv2);
+ ut_batchnorm2d_free(&b->bn1);
+ ut_batchnorm2d_free(&b->bn2);
+ if (b->has_proj) {
+ ut_conv2d_free(&b->proj);
+ ut_batchnorm2d_free(&b->projbn);
+ }
+}
+
+static ut_tensor* resnet8_forward(resnet8_t* m, ut_tensor* x, bool training, resnet8_cache_t* c) {
+ ut_tensor* s = ut_conv2d_forward(&m->stem_conv, x, c ? &c->stem_c : NULL);
+ ut_tensor* sbn = ut_batchnorm2d_forward(&m->stem_bn, s, training, c ? &c->stem_bn_c : NULL);
+ ut_free(s);
+ ut_tensor* h = ut_relu(sbn);
+ if (c) c->stem_relu_in = sbn; else ut_free(sbn);
+
+ for (int i = 0; i < 3; i++) {
+ ut_tensor* next = block_forward(&m->blocks[i], h, training, c ? &c->block_c[i] : NULL);
+ ut_free(h);
+ h = next;
+ }
+
+ int ph = h->shape.shape[2], pw = h->shape.shape[3];
+ ut_tensor* pooled = ut_global_avgpool2d(h);
+ ut_free(h);
+ if (c) c->pooled = pooled, c->pool_h = ph, c->pool_w = pw;
+ ut_tensor* logits = ut_linear_forward(&m->fc, pooled);
+ if (!c) ut_free(pooled);
+ return logits;
+}
+
+static ut_tensor* resnet8_backward(resnet8_t* m, resnet8_cache_t* c, ut_tensor* dlogits) {
+ ut_tensor* dpooled = ut_linear_backward(&m->fc, c->pooled, dlogits, m->fc_w, m->fc_b);
+ ut_tensor* dh = ut_global_avgpool2d_backward(dpooled, c->pool_h, c->pool_w);
+ ut_free(dpooled);
+
+ for (int i = 2; i >= 0; i--) {
+ ut_tensor* dprev = block_backward(&m->blocks[i], &c->block_c[i], dh);
+ ut_free(dh);
+ dh = dprev;
+ }
+
+ ut_tensor* dsbn = ut_relu_backward(dh, c->stem_relu_in);
+ ut_free(dh);
+ ut_tensor* ds = ut_batchnorm2d_backward(&m->stem_bn, &c->stem_bn_c, dsbn, m->stem_bw, m->stem_bb);
+ ut_free(dsbn);
+ ut_tensor* dx = ut_conv2d_backward(&m->stem_conv, &c->stem_c, ds, m->stem_cw, m->stem_cb);
+ ut_free(ds);
+ return dx;
+}
+
+static void resnet8_cache_free(resnet8_t* m, resnet8_cache_t* c) {
+ ut_conv2d_cache_free(&c->stem_c);
+ ut_batchnorm2d_cache_free(&c->stem_bn_c);
+ ut_free(c->stem_relu_in);
+ for (int i = 0; i < 3; i++) block_cache_free(&m->blocks[i], &c->block_c[i]);
+ ut_free(c->pooled);
+}
+
+static void resnet8_free(resnet8_t* m) {
+ ut_conv2d_free(&m->stem_conv);
+ ut_batchnorm2d_free(&m->stem_bn);
+ for (int i = 0; i < 3; i++) block_free(&m->blocks[i]);
+ ut_linear_free(&m->fc);
+}
+
+int main(void) {
+ srand(42);
+
+ ut_dev dev = UT_METAL; // flip to UT_CPU to compare against a (CPU-only) PyTorch script
+
+ printf("Loading CIFAR-10…\n");
+ const char* train_files[5] = {"cifar/data_batch_1.bin", "cifar/data_batch_2.bin",
+ "cifar/data_batch_3.bin", "cifar/data_batch_4.bin",
+ "cifar/data_batch_5.bin"};
+ const char* test_files[1] = {"cifar/test_batch.bin"};
+ cifar10_t train = cifar10_load(train_files, 5);
+ cifar10_t test = cifar10_load(test_files, 1);
+ printf("train: %d test: %d\n\n", train.n, test.n);
+
+ resnet8_t m = resnet8_alloc(dev);
+
+ ut_tensor* params[64];
+ int np = 0;
+ params[np++] = m.stem_conv.weight, params[np++] = m.stem_conv.bias;
+ params[np++] = m.stem_bn.weight, params[np++] = m.stem_bn.bias;
+ int block_off[3];
+ for (int i = 0; i < 3; i++) {
+ block_off[i] = np;
+ np += block_params(&m.blocks[i], params + np);
+ }
+ int fc_off = np;
+ params[np++] = m.fc.weight, params[np++] = m.fc.bias;
+
+ ut_adam opt = ut_adam_alloc(params, np, 1e-3f, 0.9f, 0.999f, 1e-8f, 1e-4f);
+
+ // wire each layer's gradient pointers to the optimizer's own grads[], in the
+ // same order the params[] above was built
+ m.stem_cw = opt.grads[0], m.stem_cb = opt.grads[1];
+ m.stem_bw = opt.grads[2], m.stem_bb = opt.grads[3];
+ for (int i = 0; i < 3; i++) block_wire_grads(&m.blocks[i], opt.grads + block_off[i]);
+ m.fc_w = opt.grads[fc_off], m.fc_b = opt.grads[fc_off + 1];
+
+ int B = 64, epochs = 20, batches = train.n / B;
+ ut_tensor* grad_logits = ut_alloc(2, (int[]){B, 10}, dev);
+ int* idx = malloc((size_t)train.n * sizeof(int));
+ for (int i = 0; i < train.n; i++) idx[i] = i;
+
+ for (int ep = 0; ep < epochs; ep++) {
+ cifar10_shuffle(idx, train.n);
+ clock_t t0 = clock();
+ float loss_sum = 0;
+ int correct = 0;
+
+ for (int bi = 0; bi < batches; bi++) {
+ float bx[B * CIFAR10_PIX];
+ int bl[B];
+ for (int i = 0; i < B; i++) {
+ int ii = idx[bi * B + i];
+ memcpy(bx + i * CIFAR10_PIX, train.imgs + ii * CIFAR10_PIX, CIFAR10_PIX * sizeof(float));
+ bl[i] = train.labels[ii];
+ }
+ ut_tensor* x = ut_from_data(4, (int[]){B, 3, 32, 32}, bx, dev);
+
+ resnet8_cache_t c;
+ ut_tensor* logits = resnet8_forward(&m, x, true, &c);
+ float loss = ut_cross_entropy(logits, bl, grad_logits);
+ ut_sync_cpu(logits);
+ for (int i = 0; i < B; i++)
+ if (cifar10_argmax(logits, i) == bl[i]) correct++;
+
+ ut_tensor* dx = resnet8_backward(&m, &c, grad_logits);
+ ut_adam_step(&opt, 5.0f);
+ loss_sum += loss;
+
+ resnet8_cache_free(&m, &c);
+ ut_free_all(x, logits, dx);
+ }
+
+ float secs = (float)(clock() - t0) / (float)CLOCKS_PER_SEC;
+
+ int test_ok = 0;
+ for (int i = 0; i < test.n; i += B) {
+ int nb = (i + B <= test.n) ? B : (test.n - i);
+ ut_tensor* tx = ut_from_data(4, (int[]){nb, 3, 32, 32}, test.imgs + i * CIFAR10_PIX, dev);
+ ut_tensor* tl = resnet8_forward(&m, tx, false, NULL);
+ ut_sync_cpu(tl);
+ for (int j = 0; j < nb; j++)
+ if (cifar10_argmax(tl, j) == test.labels[i + j]) test_ok++;
+ ut_free_all(tx, tl);
+ }
+
+ printf("epoch %2d loss %7.4f train %5.1f%% test %5.1f%% %6.1fs\n", ep + 1,
+ loss_sum / (float)batches, 100.f * (float)correct / (float)(batches * B),
+ 100.f * (float)test_ok / (float)test.n, secs);
+ }
+
+ free(idx);
+ ut_free(grad_logits);
+ ut_adam_free(&opt);
+ resnet8_free(&m);
+ cifar10_free(&train);
+ cifar10_free(&test);
+ return 0;
+}