← Commits · aafe8384
aafe83846cb9e25fe82b4b756248a1104a31cc20
diff --git a/test.c b/test.c
index d97187d..42a35a7 100644
--- a/test.c
+++ b/test.c
@@ -320,6 +320,62 @@ static void test_batchnorm2d_backward(ut_dev dev) {
ut_batchnorm2d_free(&bn);
}
+static void test_global_avgpool2d(ut_dev dev) {
+ // x:[1,2,2,2]: ch0=[1,2,3,4] mean=2.5, ch1=[10,20,30,40] mean=25
+ ut_tensor* x = ut_from_data(4, (int[]){1, 2, 2, 2},
+ (float[]){1.f, 2.f, 3.f, 4.f, 10.f, 20.f, 30.f, 40.f}, dev);
+ ut_tensor* out = ut_global_avgpool2d(x);
+ ut_sync_cpu(out);
+ assert(out->shape.ndim == 2 && out->shape.shape[0] == 1 && out->shape.shape[1] == 2);
+ assert_data(out, ((float[]){2.5f, 25.f}), 1e-4f);
+
+ // grad_out=[4,8] -> uniform over each channel's 4 positions: 4/4=1, 8/4=2
+ ut_tensor* go = ut_from_data(2, (int[]){1, 2}, (float[]){4.f, 8.f}, dev);
+ ut_tensor* dx = ut_global_avgpool2d_backward(go, 2, 2);
+ ut_sync_cpu(dx);
+ assert_data(dx, ((float[]){1.f, 1.f, 1.f, 1.f, 2.f, 2.f, 2.f, 2.f}), 1e-4f);
+
+ ut_free_all(x, out, go, dx);
+}
+
+static void test_maxpool2d(ut_dev dev) {
+ // x:[1,1,3,3]=[[1,2,3],[4,5,6],[7,8,9]], kh=kw=2,s=1,p=0 -> out=[5,6,8,9]
+ // (window maxes sit at input positions 4,5,7,8)
+ ut_tensor* x = ut_from_data(4, (int[]){1, 1, 3, 3},
+ (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}, dev);
+ ut_maxpool2d_cache c;
+ ut_tensor* out = ut_maxpool2d(x, 2, 2, 1, 0, &c);
+ ut_sync_cpu(out);
+ assert(out->shape.ndim == 4 && out->shape.shape[2] == 2 && out->shape.shape[3] == 2);
+ assert_data(out, ((float[]){5.f, 6.f, 8.f, 9.f}), 1e-6f);
+
+ // one-hot grad on the first output -> all gradient routed to input position 4
+ ut_tensor* go = ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
+ ut_tensor* dx = ut_maxpool2d_backward(&c, go);
+ ut_sync_cpu(dx);
+ assert_data(dx, ((float[]){0.f, 0.f, 0.f, 0.f, 1.f, 0.f, 0.f, 0.f, 0.f}), 1e-6f);
+
+ ut_free_all(x, out, go, dx);
+ ut_maxpool2d_cache_free(&c);
+}
+
+static void test_avgpool2d(ut_dev dev) {
+ // same x as maxpool; kh=kw=2,s=1,p=0 -> window means = [3,4,6,7]
+ ut_tensor* x = ut_from_data(4, (int[]){1, 1, 3, 3},
+ (float[]){1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}, dev);
+ ut_tensor* out = ut_avgpool2d(x, 2, 2, 1, 0);
+ ut_sync_cpu(out);
+ assert_data(out, ((float[]){3.f, 4.f, 6.f, 7.f}), 1e-6f);
+
+ // one-hot grad on the first output -> 0.25 spread over its 2x2 window (0,1,3,4)
+ ut_tensor* go = ut_from_data(4, (int[]){1, 1, 2, 2}, (float[]){1.f, 0.f, 0.f, 0.f}, dev);
+ ut_tensor* dx = ut_avgpool2d_backward(go, 1, 1, 3, 3, 2, 2, 1, 0);
+ ut_sync_cpu(dx);
+ assert_data(dx, ((float[]){0.25f, 0.25f, 0.f, 0.25f, 0.25f, 0.f, 0.f, 0.f, 0.f}), 1e-6f);
+
+ ut_free_all(x, out, go, dx);
+}
+
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);
@@ -546,6 +602,13 @@ int main() {
test_batchnorm2d_backward(UT_CPU);
test_batchnorm2d_backward(UT_METAL);
+ test_global_avgpool2d(UT_CPU);
+ test_global_avgpool2d(UT_METAL);
+ test_maxpool2d(UT_CPU);
+ test_maxpool2d(UT_METAL);
+ test_avgpool2d(UT_CPU);
+ test_avgpool2d(UT_METAL);
+
test_im2col(UT_CPU);
test_im2col(UT_METAL);
test_col2im(UT_CPU);
diff --git a/utensil.h b/utensil.h
index e61b209..b88604e 100644
--- a/utensil.h
+++ b/utensil.h
@@ -83,6 +83,12 @@ typedef struct ut_batchnorm2d_cache {
ut_tensor* rstd; // per-channel reciprocal std-dev
} ut_batchnorm2d_cache;
+typedef struct ut_maxpool2d_cache {
+ int* argmax; // flat input index of the max, per output element (plain array — CPU-only
+ // bookkeeping, never touches the GPU)
+ int n, c, h, w; // input shape, needed to size dx in backward
+} ut_maxpool2d_cache;
+
typedef struct ut_sgd {
ut_tensor** params; // pointers to model parameters (not owned)
ut_tensor** grads; // gradient accumulators (owned)
@@ -1617,6 +1623,133 @@ void ut_batchnorm2d_free(ut_batchnorm2d* l) {
l->weight = l->bias = l->running_mean = l->running_var = NULL;
}
+// =========================================================
+// Pooling
+// =========================================================
+
+// x: [N,C,H,W] -> [N,C], averaging each channel over H,W. No cache: backward
+// only needs H,W, which the caller already has from x's own shape.
+ut_tensor* ut_global_avgpool2d(ut_tensor* x) {
+ int N = x->shape.shape[0], C = x->shape.shape[1], HW = x->shape.shape[2] * x->shape.shape[3];
+ ut_sync_cpu(x);
+ ut_tensor* out = ut_alloc(2, (int[]){N, C}, x->dev);
+ for (int n = 0; n < N; n++)
+ for (int c = 0; c < C; c++) {
+ float sum = 0;
+ for (int i = 0; i < HW; i++) sum += x->data[(n * C + c) * HW + i];
+ out->data[n * C + c] = sum / (float)HW;
+ }
+ out->dirty_gpu = true;
+ return out;
+}
+
+ut_tensor* ut_global_avgpool2d_backward(ut_tensor* grad_out, int H, int W) {
+ int N = grad_out->shape.shape[0], C = grad_out->shape.shape[1], HW = H * W;
+ ut_sync_cpu(grad_out);
+ ut_tensor* dx = ut_alloc(4, (int[]){N, C, H, W}, grad_out->dev);
+ for (int n = 0; n < N; n++)
+ for (int c = 0; c < C; c++) {
+ float g = grad_out->data[n * C + c] / (float)HW;
+ for (int i = 0; i < HW; i++) dx->data[(n * C + c) * HW + i] = g;
+ }
+ dx->dirty_gpu = true;
+ return dx;
+}
+
+// x: [N,C,H,W] -> [N,C,Ho,Wo], max over each kh x kw window per channel
+ut_tensor* ut_maxpool2d(ut_tensor* x, int kh, int kw, int stride, int pad,
+ ut_maxpool2d_cache* cache) {
+ int N = x->shape.shape[0], C = x->shape.shape[1], H = x->shape.shape[2], W = x->shape.shape[3];
+ int Ho = (H + 2 * pad - kh) / stride + 1, Wo = (W + 2 * pad - kw) / stride + 1;
+ ut_sync_cpu(x);
+ ut_tensor* out = ut_alloc(4, (int[]){N, C, Ho, Wo}, x->dev);
+ int* argmax = cache ? malloc((size_t)N * C * Ho * Wo * sizeof(int)) : NULL;
+ for (int n = 0; n < N; n++)
+ for (int c = 0; c < C; c++)
+ for (int oh = 0; oh < Ho; oh++)
+ for (int ow = 0; ow < Wo; ow++) {
+ float best = -FLT_MAX;
+ int best_idx = -1;
+ for (int hh = 0; hh < kh; hh++)
+ for (int ww = 0; ww < kw; ww++) {
+ int ih = oh * stride - pad + hh, iw = ow * stride - pad + ww;
+ if (ih < 0 || ih >= H || iw < 0 || iw >= W) continue;
+ int idx = ((n * C + c) * H + ih) * W + iw;
+ if (x->data[idx] > best) best = x->data[idx], best_idx = idx;
+ }
+ int oidx = ((n * C + c) * Ho + oh) * Wo + ow;
+ out->data[oidx] = best;
+ if (cache) argmax[oidx] = best_idx;
+ }
+ out->dirty_gpu = true;
+ if (cache) {
+ cache->argmax = argmax;
+ cache->n = N, cache->c = C, cache->h = H, cache->w = W;
+ }
+ return out;
+}
+
+ut_tensor* ut_maxpool2d_backward(ut_maxpool2d_cache* cache, ut_tensor* grad_out) {
+ ut_sync_cpu(grad_out);
+ ut_tensor* dx = ut_alloc(4, (int[]){cache->n, cache->c, cache->h, cache->w}, grad_out->dev);
+ memset(dx->data, 0, (size_t)dx->shape.nelem * sizeof(float));
+ for (int i = 0; i < grad_out->shape.nelem; i++) dx->data[cache->argmax[i]] += grad_out->data[i];
+ dx->dirty_gpu = true;
+ return dx;
+}
+
+void ut_maxpool2d_cache_free(ut_maxpool2d_cache* c) {
+ free(c->argmax);
+ c->argmax = NULL;
+}
+
+// x: [N,C,H,W] -> [N,C,Ho,Wo], mean over each kh x kw window per channel
+// (padding counts as zero, matching PyTorch's count_include_pad=True default)
+ut_tensor* ut_avgpool2d(ut_tensor* x, int kh, int kw, int stride, int pad) {
+ int N = x->shape.shape[0], C = x->shape.shape[1], H = x->shape.shape[2], W = x->shape.shape[3];
+ int Ho = (H + 2 * pad - kh) / stride + 1, Wo = (W + 2 * pad - kw) / stride + 1;
+ ut_sync_cpu(x);
+ ut_tensor* out = ut_alloc(4, (int[]){N, C, Ho, Wo}, x->dev);
+ for (int n = 0; n < N; n++)
+ for (int c = 0; c < C; c++)
+ for (int oh = 0; oh < Ho; oh++)
+ for (int ow = 0; ow < Wo; ow++) {
+ float sum = 0;
+ for (int hh = 0; hh < kh; hh++)
+ for (int ww = 0; ww < kw; ww++) {
+ int ih = oh * stride - pad + hh, iw = ow * stride - pad + ww;
+ if (ih >= 0 && ih < H && iw >= 0 && iw < W)
+ sum += x->data[((n * C + c) * H + ih) * W + iw];
+ }
+ out->data[((n * C + c) * Ho + oh) * Wo + ow] = sum / (float)(kh * kw);
+ }
+ out->dirty_gpu = true;
+ return out;
+}
+
+ut_tensor* ut_avgpool2d_backward(ut_tensor* grad_out, int N, int C, int H, int W, int kh, int kw,
+ int stride, int pad) {
+ int Ho = (H + 2 * pad - kh) / stride + 1, Wo = (W + 2 * pad - kw) / stride + 1;
+ ut_sync_cpu(grad_out);
+ ut_tensor* dx = ut_alloc(4, (int[]){N, C, H, W}, grad_out->dev);
+ memset(dx->data, 0, (size_t)dx->shape.nelem * sizeof(float));
+ float scale = 1.f / (float)(kh * kw);
+ for (int n = 0; n < N; n++)
+ for (int c = 0; c < C; c++)
+ for (int oh = 0; oh < Ho; oh++)
+ for (int ow = 0; ow < Wo; ow++) {
+ float g = grad_out->data[((n * C + c) * Ho + oh) * Wo + ow] * scale;
+ for (int hh = 0; hh < kh; hh++)
+ for (int ww = 0; ww < kw; ww++) {
+ int ih = oh * stride - pad + hh, iw = ow * stride - pad + ww;
+ if (ih >= 0 && ih < H && iw >= 0 && iw < W)
+ dx->data[((n * C + c) * H + ih) * W + iw] += g;
+ }
+ }
+ dx->dirty_gpu = true;
+ return dx;
+}
+
// =========================================================
// Softmax
// =========================================================