Merge pull request #6 from skainet-dev/codex/add-maxpool2d-support

Add MaxPool2d layer and DSL

Author: skainet-dev

core/src/commonMain/kotlin/sk/ai/net/dsl/NetworkBuilder.kt

@@ -7,6 +7,7 @@ import sk.ai.net.nn.Flatten
 import sk.ai.net.nn.Input
 import sk.ai.net.nn.Linear
 import sk.ai.net.nn.Conv2d
+import sk.ai.net.nn.MaxPool2d
 import sk.ai.net.nn.Module
 import sk.ai.net.nn.topology.MLP
 
@@ -31,6 +32,8 @@ interface NeuralNetworkDsl : NetworkDslItem {
 
     fun conv2d(id: String = "", content: CONV2D.() -> Unit = {})
 
+    fun maxPool2d(id: String = "", content: MAXPOOL2D.() -> Unit = {})
+
     fun dense(outputDimension: Int, id: String = "", content: DENSE.() -> Unit = {})
 }
 
@@ -55,6 +58,12 @@ interface CONV2D : NetworkDslItem {
     var padding: Int
 }
 
+@NetworkDsl
+interface MAXPOOL2D : NetworkDslItem {
+    var kernelSize: Int
+    var stride: Int
+}
+
 
 private fun getDefaultName(id: String, s: String, size: Int): String {
     if (id.isNotEmpty()) return id
@@ -157,6 +166,18 @@ class Conv2dImpl(
     )
 }
 
+class MaxPool2dImpl(
+    override var kernelSize: Int = 2,
+    override var stride: Int = 2,
+    private val id: String
+) : MAXPOOL2D {
+    fun create(): Module = MaxPool2d(
+        kernelSize = kernelSize,
+        stride = stride,
+        name = id
+    )
+}
+
 private class NeuralNetworkDslImpl : NeuralNetworkDsl {
 
     val modules = mutableListOf<Module>()
@@ -186,6 +207,14 @@ private class NeuralNetworkDslImpl : NeuralNetworkDsl {
         modules += impl.create()
     }
 
+    override fun maxPool2d(id: String, content: MAXPOOL2D.() -> Unit) {
+        val impl = MaxPool2dImpl(
+            id = getDefaultName(id, "maxPool2d", modules.size)
+        )
+        impl.content()
+        modules += impl.create()
+    }
+
     override fun dense(outputDimension: Int, id: String, content: DENSE.() -> Unit) {
         val inputDimension = lastDimension
         lastDimension = outputDimension

core/src/commonMain/kotlin/sk/ai/net/nn/MaxPool2d.kt

@@ -0,0 +1,72 @@
+package sk.ai.net.nn
+
+import sk.ai.net.Shape
+import sk.ai.net.Tensor
+import sk.ai.net.impl.DoublesTensor
+
+/**
+ * 2D max pooling layer.
+ * Works with tensors of shape (N, C, H, W) or (C, H, W).
+ */
+class MaxPool2d(
+    val kernelSize: Int,
+    val stride: Int = kernelSize,
+    override val name: String = "MaxPool2d"
+) : Module() {
+    override val modules: List<Module>
+        get() = emptyList()
+
+    override fun forward(input: Tensor): Tensor = maxPool2d(input)
+
+    private fun maxPool2d(input: Tensor): Tensor {
+        val tensor = input as DoublesTensor
+        val shape = tensor.shape
+        require(shape.rank == 3 || shape.rank == 4) {
+            "MaxPool2d expected 3D or 4D input tensor, but got shape $shape"
+        }
+        val batchSize: Int
+        val channels: Int
+        val height: Int
+        val width: Int
+        if (shape.rank == 4) {
+            batchSize = shape[0]
+            channels = shape[1]
+            height = shape[2]
+            width = shape[3]
+        } else {
+            batchSize = 1
+            channels = shape[0]
+            height = shape[1]
+            width = shape[2]
+        }
+
+        val outH = (height - kernelSize) / stride + 1
+        val outW = (width - kernelSize) / stride + 1
+        val outElements = DoubleArray(batchSize * channels * outH * outW)
+        var idx = 0
+        for (n in 0 until batchSize) {
+            for (c in 0 until channels) {
+                for (i in 0 until outH) {
+                    for (j in 0 until outW) {
+                        var maxVal = Double.NEGATIVE_INFINITY
+                        for (ki in 0 until kernelSize) {
+                            for (kj in 0 until kernelSize) {
+                                val h = i * stride + ki
+                                val w = j * stride + kj
+                                val value = if (shape.rank == 4) {
+                                    tensor[n, c, h, w]
+                                } else {
+                                    tensor[c, h, w]
+                                }
+                                if (value > maxVal) maxVal = value
+                            }
+                        }
+                        outElements[idx++] = maxVal
+                    }
+                }
+            }
+        }
+        val outShape = Shape(batchSize, channels, outH, outW)
+        return DoublesTensor(outShape, outElements)
+    }
+}

core/src/commonTest/kotlin/sk/ai/net/nn/MaxPool2dTest.kt

@@ -0,0 +1,44 @@
+package sk.ai.net.nn
+
+import sk.ai.net.Shape
+import sk.ai.net.dsl.network
+import sk.ai.net.impl.DoublesTensor
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+
+class MaxPool2dTest {
+    @Test
+    fun max_pool2d_basic() {
+        val input = DoublesTensor(
+            Shape(1, 1, 4, 4),
+            doubleArrayOf(
+                1.0, 2.0, 3.0, 4.0,
+                5.0, 6.0, 7.0, 8.0,
+                9.0, 10.0, 11.0, 12.0,
+                13.0, 14.0, 15.0, 16.0
+            )
+        )
+        val pool = MaxPool2d(kernelSize = 2, stride = 2)
+        val result = pool.forward(input) as DoublesTensor
+        assertEquals(Shape(1, 1, 2, 2), result.shape)
+        assertContentEquals(doubleArrayOf(6.0, 8.0, 14.0, 16.0), result.elements)
+    }
+
+    @Test
+    fun dsl_support() {
+        val module = network {
+            input(1)
+            maxPool2d {
+                kernelSize = 2
+                stride = 2
+            }
+        }
+        val mlp = module as sk.ai.net.nn.topology.MLP
+        assertTrue(mlp.modules[1] is MaxPool2d)
+        val mp = mlp.modules[1] as MaxPool2d
+        assertEquals(2, mp.kernelSize)
+        assertEquals(2, mp.stride)
+    }
+}