移植bot4cj

bot4cj/src/math/pooling.cj

@@ -0,0 +1,418 @@
+// pooling_layer.cj
+package bot4cj.math
+
+import std.math.*
+import std.collection.*
+
+/**
+ * 池化层类
+ * 设计风格与 Nn1 和 ConvLayer 保持一致
+ */
+public class PoolingLayer {
+
+    // ==================== 层配置参数 ====================
+    let nChannels: Int64
+    let poolHeight: Int64
+    let poolWidth: Int64
+    let stride: Int64
+    let padding: Int64
+    let poolType: String
+    let f: Af
+    let learningRateWeight: FP32
+    let learningRateBias: FP32
+
+    // ==================== 可学习参数 ====================
+    var weights: Matrix1
+    var biases: Matrix1
+
+    // ==================== 运行时缓存 ====================
+    var prevInput: Option<Matrix1> = Option.None
+    var prevInputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
+    var prevCombination: Option<Matrix1> = Option.None
+    var prevOutput: Option<Matrix1> = Option.None
+    var prevOutputShape: Option<(Int64, Int64, Int64, Int64)> = Option.None
+    var prevMaxIndices: Option<Matrix1> = Option.None
+
+    var weightsGrad: Option<Matrix1> = Option.None
+    var biasGrad: Option<Matrix1> = Option.None
+
+    // ==================== 构造函数 ====================
+    public init(
+        nChannels: Int64,
+        poolSize: (Int64, Int64),
+        stride: Int64,
+        padding: Int64,
+        poolType: String,
+        activation: Af,
+        lrWeight: FP32,
+        lrBias: FP32
+    ) {
+        this.nChannels = nChannels
+        poolHeight = poolSize[0]
+        poolWidth = poolSize[1]
+        this.stride = stride
+        this.padding = padding
+        this.poolType = poolType
+        this.f = activation
+        learningRateWeight = lrWeight
+        learningRateBias = lrBias
+
+        // 初始化权值
+        weights = Matrix1(1, nChannels)
+        if (poolType == "average") {
+            let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
+            for (c in 0..nChannels) {
+                weights.set(0, c, avgWeight)
+            }
+        } else {
+            for (c in 0..nChannels) {
+                weights.set(0, c, FP32(1.0))
+            }
+        }
+
+        // 初始化偏置为 0
+        biases = Matrix1(1, nChannels)
+        for (c in 0..nChannels) {
+            biases.set(0, c, FP32(0.0))
+        }
+    }
+
+    public init(
+        nChannels: Int64,
+        poolSize: (Int64, Int64),
+        poolType: String,
+        activation: Af
+    ) {
+        this.nChannels = nChannels
+        poolHeight = poolSize[0]
+        poolWidth = poolSize[1]
+        stride = poolHeight
+        padding = 0
+        this.poolType = poolType
+        this.f = activation
+        learningRateWeight = FP32(0.1)
+        learningRateBias = FP32(0.1)
+
+        weights = Matrix1(1, nChannels)
+        if (poolType == "average") {
+            let avgWeight = FP32(1.0) / FP32(Float32(poolHeight * poolWidth))
+            for (c in 0..nChannels) {
+                weights.set(0, c, avgWeight)
+            }
+        } else {
+            for (c in 0..nChannels) {
+                weights.set(0, c, FP32(1.0))
+            }
+        }
+
+        biases = Matrix1(1, nChannels)
+        for (c in 0..nChannels) {
+            biases.set(0, c, FP32(0.0))
+        }
+    }
+
+    // ==================== 前向传播 ====================
+    public func forward(x: Matrix1, inputShape: (Int64, Int64, Int64, Int64)): Matrix1 {
+        prevInput = Option.Some(x)
+        prevInputShape = Option.Some(inputShape)
+
+        let batch = inputShape[0]
+        let inH = inputShape[1]
+        let inW = inputShape[2]
+
+        let outH = (inH + 2 * padding - poolHeight) / stride + 1
+        let outW = (inW + 2 * padding - poolWidth) / stride + 1
+        let outSize = batch * outH * outW
+
+        var combination = Matrix1(outSize, nChannels)
+        var output = Matrix1(outSize, nChannels)
+        var maxIndices = Matrix1(outSize, nChannels)
+
+        for (b in 0..batch) {
+            for (oh in 0..outH) {
+                for (ow in 0..outW) {
+                    let outRow = b * outH * outW + oh * outW + ow
+
+                    for (ch in 0..nChannels) {
+                        var pooledVal = FP32(0.0)
+                        var maxIdx = 0
+
+                        if (poolType == "max") {
+                            var maxVal = FP32(-1.0e30)
+                            var kernelIdx = 0
+                            for (kh in 0..poolHeight) {
+                                for (kw in 0..poolWidth) {
+                                    let ih = oh * stride + kh - padding
+                                    let iw = ow * stride + kw - padding
+
+                                    if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                                        let inRow = b * inH * inW + ih * inW + iw
+                                        let val = x.get(inRow, ch).getOrThrow()
+                                        if (val.getValue() > maxVal.getValue()) {
+                                            maxVal = val
+                                            maxIdx = kernelIdx
+                                        }
+                                    }
+                                    kernelIdx += 1
+                                }
+                            }
+                            pooledVal = maxVal
+                        } else {
+                            var sum = FP32(0.0)
+                            var count = 0
+                            for (kh in 0..poolHeight) {
+                                for (kw in 0..poolWidth) {
+                                    let ih = oh * stride + kh - padding
+                                    let iw = ow * stride + kw - padding
+
+                                    if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                                        let inRow = b * inH * inW + ih * inW + iw
+                                        sum = sum + x.get(inRow, ch).getOrThrow()
+                                        count += 1
+                                    }
+                                }
+                            }
+                            if (count > 0) {
+                                pooledVal = sum / FP32(Float32(count))
+                            }
+                        }
+
+                        combination.set(outRow, ch, pooledVal)
+                        maxIndices.set(outRow, ch, FP32(Float32(maxIdx)))
+                    }
+                }
+            }
+        }
+
+        // 逐通道应用权值和偏置
+        for (r in 0..outSize) {
+            for (ch in 0..nChannels) {
+                let val = combination.get(r, ch).getOrThrow()
+                let w = weights.get(0, ch).getOrThrow()
+                let b = biases.get(0, ch).getOrThrow()
+                combination.set(r, ch, val * w + b)
+            }
+        }
+
+        prevCombination = Option.Some(combination)
+
+        output = f.forward(combination)
+        prevOutput = Option.Some(output)
+        prevOutputShape = Option.Some((batch, outH, outW, nChannels))
+        prevMaxIndices = Option.Some(maxIndices)
+
+        return output
+    }
+
+    // ==================== 反向传播 ====================
+    public func backward(errs: Matrix1): Matrix1 {
+        var x = Matrix1(1, 1)
+        match (prevInput) {
+            case Some(val) => x = val
+            case _ => throw Exception("No input data")
+        }
+        
+        var shape = (0, 0, 0, 0)
+        match (prevInputShape) {
+            case Some(val) => shape = val
+            case _ => throw Exception("No input shape")
+        }
+        
+        var combination = Matrix1(1, 1)
+        match (prevCombination) {
+            case Some(val) => combination = val
+            case _ => throw Exception("No combination")
+        }
+        
+        var outShape = (0, 0, 0, 0)
+        match (prevOutputShape) {
+            case Some(val) => outShape = val
+            case _ => throw Exception("No output shape")
+        }
+
+        var maxIndices = Matrix1(1, 1)
+        let hasMaxIndices = match (prevMaxIndices) {
+            case Some(_) => true
+            case _ => false
+        }
+        if (hasMaxIndices) {
+            maxIndices = match (prevMaxIndices) {
+                case Some(val) => val
+                case _ => Matrix1(1, 1)
+            }
+        }
+        
+        let batch = shape[0]
+        let inH = shape[1]
+        let inW = shape[2]
+        let outH = outShape[1]
+        let outW = outShape[2]
+        
+        let delta = errs * f.backward(combination)
+
+        calcWeightBiasGradient(delta, combination)
+
+        let errsUp = calcUpsampledError(delta, maxIndices, batch, inH, inW, outH, outW)
+        
+        return errsUp
+    }
+
+    /**
+     * 计算权值和偏置梯度
+     */
+    private func calcWeightBiasGradient(delta: Matrix1, combination: Matrix1): Unit {
+        let outSize = delta.get_row()
+        
+        var wGrad = Matrix1(1, nChannels)
+        var bGrad = Matrix1(1, nChannels)
+
+        for (ch in 0..nChannels) {
+            var wSum = FP32(0.0)
+            var bSum = FP32(0.0)
+            for (r in 0..outSize) {
+                let d = delta.get(r, ch).getOrThrow()
+                bSum = bSum + d
+                let comb = combination.get(r, ch).getOrThrow()
+                wSum = wSum + d * comb
+            }
+            wGrad.set(0, ch, wSum)
+            bGrad.set(0, ch, bSum)
+        }
+
+        weightsGrad = Option.Some(wGrad)
+        biasGrad = Option.Some(bGrad)
+    }
+
+    /**
+     * 上采样误差（池化层的反向传播核心）
+     */
+    private func calcUpsampledError(
+        delta: Matrix1,
+        maxIndices: Matrix1,
+        batch: Int64,
+        inH: Int64,
+        inW: Int64,
+        outH: Int64,
+        outW: Int64
+    ): Matrix1 {
+        let inSize = batch * inH * inW
+        var errsUp = Matrix1(inSize, nChannels)
+
+        for (b in 0..batch) {
+            for (oh in 0..outH) {
+                for (ow in 0..outW) {
+                    let deltaRow = b * outH * outW + oh * outW + ow
+
+                    for (ch in 0..nChannels) {
+                        let deltaVal = delta.get(deltaRow, ch).getOrThrow()
+                        let w = weights.get(0, ch).getOrThrow()
+
+                        if (poolType == "max") {
+                            let maxIdx = Int64(maxIndices.get(deltaRow, ch).getOrThrow().getValue())
+                            
+                            var kernelIdx = 0
+                            var found = false
+                            for (kh in 0..poolHeight) {
+                                for (kw in 0..poolWidth) {
+                                    if (kernelIdx == maxIdx && !found) {
+                                        let ih = oh * stride + kh - padding
+                                        let iw = ow * stride + kw - padding
+
+                                        if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                                            let inRow = b * inH * inW + ih * inW + iw
+                                            errsUp.set(inRow, ch, deltaVal * w)
+                                            found = true
+                                        }
+                                    }
+                                    kernelIdx += 1
+                                }
+                            }
+                        } else {
+                            var count = 0
+                            for (kh in 0..poolHeight) {
+                                for (kw in 0..poolWidth) {
+                                    let ih = oh * stride + kh - padding
+                                    let iw = ow * stride + kw - padding
+                                    if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                                        count += 1
+                                    }
+                                }
+                            }
+                            
+                            if (count > 0) {
+                                let avgDelta = deltaVal * w / FP32(Float32(count))
+                                for (kh in 0..poolHeight) {
+                                    for (kw in 0..poolWidth) {
+                                        let ih = oh * stride + kh - padding
+                                        let iw = ow * stride + kw - padding
+                                        if (ih >= 0 && ih < inH && iw >= 0 && iw < inW) {
+                                            let inRow = b * inH * inW + ih * inW + iw
+                                            let oldVal = errsUp.get(inRow, ch).getOrThrow()
+                                            errsUp.set(inRow, ch, oldVal + avgDelta)
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        
+        return errsUp
+    }
+
+    // ==================== 参数更新 ====================
+    public func update(discount: Float32): Unit {
+        var wGrad = Matrix1(1, 1)
+        match (weightsGrad) {
+            case Some(val) => wGrad = val
+            case _ => throw Exception("No weight gradients")
+        }
+        
+        var bGrad = Matrix1(1, 1)
+        match (biasGrad) {
+            case Some(val) => bGrad = val
+            case _ => throw Exception("No bias gradients")
+        }
+        
+        for (ch in 0..nChannels) {
+            let oldW = weights.get(0, ch).getOrThrow()
+            let gW = wGrad.get(0, ch).getOrThrow()
+            weights.set(0, ch, oldW + gW * learningRateWeight * FP32(discount))
+        }
+        
+        for (ch in 0..nChannels) {
+            let oldB = biases.get(0, ch).getOrThrow()
+            let gB = bGrad.get(0, ch).getOrThrow()
+            biases.set(0, ch, oldB + gB * learningRateBias * FP32(discount))
+        }
+    }
+
+    // ==================== 形状访问器 ====================
+    public func getInShape(): (Int64, Int64) {
+        (nChannels, poolHeight * poolWidth)
+    }
+
+    public func getOutShape(): Int64 {
+        nChannels
+    }
+
+    public func getPoolType(): String {
+        poolType
+    }
+
+    public func getLR(): FP32 {
+        learningRateWeight
+    }
+
+    public func getFType(): String {
+        f.get_type()
+    }
+
+    public func calculateOutputSize(inputHeight: Int64, inputWidth: Int64): (Int64, Int64) {
+        let outH = (inputHeight + 2 * padding - poolHeight) / stride + 1
+        let outW = (inputWidth + 2 * padding - poolWidth) / stride + 1
+        (outH, outW)
+    }
+}