✨ feat(layer_2d): DecorelateRGB

Docs/Contributing/CHANGELOG.md

@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+🪜 **layer_2d**: DecorelateRGB ([#23](https://github.com/owkin/MAKit/pull/23))\
 🪜 **layer_2d**: RDFT2Image ([#22](https://github.com/owkin/MAKit/pull/22))\
 🪜 **core:** Sigmoid activation ([#21](https://github.com/owkin/MAKit/pull/21))\
 👷 **ci:** remove swift action ([#20](https://github.com/owkin/MAKit/pull/20))\

Sources/MAKit/Layer2D/Activation2D.swift

@@ -289,7 +289,6 @@ public class Activation2D: Layer2D
                     }
                 }}}
             }
-
             propagateDirty()
         }
     }

Sources/MAKit/Layer2D/AdaptiveAvgPool2D.swift

@@ -733,7 +733,6 @@ public class AdaptiveAvgPool2D: Layer2D
                     }}
                 }}
             }
-
             propagateDirty()
         }
     }

Sources/MAKit/Layer2D/BN2D.swift

@@ -612,7 +612,6 @@ public class BN2D: Activation2D, LayerUpdate, LayerExtract
                     }
                 }}}
             }
-
             propagateDirty()
         }
     }

Sources/MAKit/Layer2D/DecorelateRGB.swift

@@ -0,0 +1,315 @@
+//
+// DecorelateRGB.swift
+// MAKit
+//
+// Created by Jean-François Reboud on 26/11/2022.
+//
+
+///
+/// Layer with a 2D shape neural structure.
+///
+/// This layer decorelates RGB color of an image.
+///
+public class DecorelateRGB: Layer2D
+{
+    let _correlation: [Double]
+
+    private enum Keys: String, CodingKey
+    {
+        case correlation
+    }
+
+    ///
+    /// Create a layer with a 2D shape neural structure.
+    ///
+    /// - Parameters:
+    ///     - layerPrev: Previous layer that has been queued to the model.
+    ///     - params: Contextual parameters linking to the model.
+    ///
+    public init(layerPrev: Layer2D,
+                correlation: [Double],
+                params: MAKit.Model.Params)
+    {
+        _correlation = correlation
+
+        let width = layerPrev.width
+        let height = layerPrev.height
+        let nbChannels = layerPrev.nbChannels
+
+        if nbChannels != 3
+        {
+            fatalError("DecorelateRGB can only be used with 3 channels.")
+        }
+        super.init(layerPrev: layerPrev,
+                   nbChannels: nbChannels,
+                   height: height,
+                   width: width,
+                   params: params)
+    }
+
+    ///
+    /// Decode from the disk.
+    ///
+    /// Throw an error if reading from the decoder fails, or
+    /// if the data read is corrupted or otherwise invalid.
+    ///
+    /// - Parameter decoder: The decoder to read data from.
+    ///
+    public required init(from decoder: Decoder) throws
+    {
+        let values = try decoder.container(keyedBy: Keys.self)
+        _correlation = try values.decode(
+            [Double].self, forKey: Keys.correlation
+        )
+        try super.init(from: decoder)
+    }
+
+    ///
+    /// Encode to the disk.
+    ///
+    /// If the value fails to encode anything, `encoder` will encode an empty
+    /// keyed container in its place.
+    ///
+    /// Throw an error if any values are invalid for the given
+    /// encoder's format.
+    ///
+    /// - Parameter encoder: The encoder to write data to.
+    ///
+    public override func encode(to encoder: Encoder) throws
+    {
+        var container = encoder.container(keyedBy: Keys.self)
+        try container.encode(_correlation, forKey: Keys.correlation)
+        try super.encode(to: encoder)
+    }
+
+    ///
+    /// Create a layer with same values as this.
+    ///
+    /// - Parameters:
+    ///     - mapping: Dictionary allowing to find the layer associated to some id.
+    ///     This dictionary is particularly useful when the different layers cannot access
+    ///     their `layerPrev`.
+    ///     - inPlace: Whether hard resources should be copied as is.
+    ///
+    /// - Returns: A new layer. When `inPlace` is false, `initKernel` is
+    /// necessary in order to recreate hard resources.
+    ///
+    public override func copy(
+        mapping: Dictionary<Int, Layer>,
+        inPlace: Bool) -> Layer
+    {
+        let context = ModelContext(name: "", curID: 0)
+        let layerPrev = mapping[idPrev] as! Layer2D
+
+        let params = MAKit.Model.Params(context: context)
+        params.context.curID = id
+
+        let layer = DecorelateRGB(
+            layerPrev: layerPrev,
+            correlation: _correlation,
+            params: params
+        )
+        return layer
+    }
+
+    ///
+    /// Apply the forward pass of the Gradient Checking in CPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGCCPU() throws
+    {
+        if let layerPrev = self.layerPrev as? Layer2D
+        {
+            try checkStateCPU(batchSize: batchSize)
+
+            let nbGC = layerPrev.nbGC
+            for depth in 0..<nbChannels
+            {
+                for i in 0..<height {
+                for j in 0..<width
+                {
+                    neurons[depth].get(i, j)!.initGC(
+                        batchSize: batchSize,
+                        nbGC: nbGC
+                    )
+                }}
+            }
+
+            let neuronsPrev = layerPrev.neurons
+            for batch in 0..<batchSize {
+            for elem in 0..<nbGC {
+            for depth in 0..<nbChannels
+            {
+                let block = depth / 3
+                let res = depth % 3
+
+                for i in 0..<height {
+                for j in 0..<width
+                {
+                    var sum: Double = 0.0
+                    for k in 0..<3
+                    {
+                        sum +=
+                            neuronsPrev[block * 3 + k].get(i, j)!
+                                .gc[batch][elem].out *
+                            _correlation[res * 3 + k]
+                    }
+                    neurons[depth].get(i, j)!.gc[batch][elem].out = sum
+                }}
+            }}}
+        }
+    }
+
+    ///
+    /// Apply the forward pass of the Gradient Checking in GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGCGPU() throws
+    {
+        try forwardGCCPU()
+    }
+
+    ///
+    /// Apply the forward pass in the CPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardCPU() throws
+    {
+        if let layerPrev = self.layerPrev as? Layer2D
+        {
+            try checkStateCPU(batchSize: batchSize)
+
+            let neuronsPrev = layerPrev.neurons
+            for elem in 0..<batchSize {
+            for depth in 0..<nbChannels
+            {
+                let block = depth / 3
+                let res = depth % 3
+
+                for i in 0..<height {
+                for j in 0..<width
+                {
+                    var sum: Double = 0.0
+                    for k in 0..<3
+                    {
+                        sum +=
+                            neuronsPrev[block * 3 + k].get(i, j)!.v[elem].out *
+                            _correlation[res * 3 + k]
+                    }
+                    neurons[depth].get(i, j)!.v[elem].out = sum
+                }}
+            }}
+        }
+    }
+
+    ///
+    /// Apply the forward pass in the GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func forwardGPU() throws
+    {
+        if let layerPrev = self.layerPrev as? Layer2D
+        {
+            try checkStateForwardGPU(batchSize: batchSize)
+
+            let pNbChannels: [UInt32] = [UInt32(nbChannels)]
+            let pNbBatch: [UInt32] = [UInt32(batchSize)]
+            let pDimensions: [UInt32] = [UInt32(width), UInt32(height)]
+            let correlation: [Float] = _correlation.map { Float($0) }
+
+            let command = MetalKernel.get.createCommand(
+                "decorelateRGBForward", deviceID: deviceID
+            )
+            command.setBuffer(layerPrev.outs.metal, atIndex: 0)
+            command.setBytes(correlation, atIndex: 1)
+            command.setBytes(pNbChannels, atIndex: 2)
+            command.setBytes(pDimensions, atIndex: 3)
+            command.setBytes(pNbBatch, atIndex: 4)
+            command.setBuffer(outs.metal, atIndex: 5)
+
+            command.dispatchThreads(
+                width: width * nbChannels,
+                height: height * batchSize
+            )
+            command.enqueue()
+        }
+    }
+
+    /// Apply the backward pass in the CPU execution context.
+    public override func backwardCPU()
+    {
+        if let layerPrev = self.layerPrev as? Layer2D, mustComputeBackward
+        {
+            let neuronsPrev = layerPrev.neurons
+            for elem in 0..<batchSize {
+            for depth in 0..<nbChannels
+            {
+                let block = depth / 3
+                let res = depth % 3
+
+                for i in 0..<height {
+                for j in 0..<width
+                {
+                    var sum: Double = 0.0
+                    for k in 0..<3
+                    {
+                        sum += _correlation[k * 3 + res] *
+                            neurons[block * 3 + k].get(i, j)!.v[elem].delta
+                    }
+
+                    if layerPrev.dirty
+                    {
+                        neuronsPrev[depth].get(i, j)!.v[elem].delta = sum
+                    }
+                    else
+                    {
+                        neuronsPrev[depth].get(i, j)!.v[elem].delta += sum
+                    }
+                }}
+            }}
+            propagateDirty()
+        }
+    }
+
+    ///
+    /// Apply the backward pass in the GPU execution context.
+    ///
+    /// Throw an error if batch size is greater than the first batch size.
+    ///
+    public override func backwardGPU() throws
+    {
+        if let layerPrev = self.layerPrev as? Layer2D, mustComputeBackward
+        {
+            try layerPrev.checkStateBackwardGPU(batchSize: batchSize)
+
+            let pNbChannels: [UInt32] = [UInt32(nbChannels)]
+            let pNbBatch: [UInt32] = [UInt32(batchSize)]
+            let pDimensions: [UInt32] = [UInt32(width), UInt32(height)]
+            let pDirty: [UInt32] = layerPrev.dirty ? [1] : [0]
+            let correlation: [Float] = _correlation.map { Float($0) }
+
+            let command = MetalKernel.get.createCommand(
+                "decorelateRGBBackward", deviceID: deviceID
+            )
+            command.setBuffer(delta.metal, atIndex: 0)
+            command.setBytes(correlation, atIndex: 1)
+            command.setBytes(pNbChannels, atIndex: 2)
+            command.setBytes(pDimensions, atIndex: 3)
+            command.setBytes(pNbBatch, atIndex: 4)
+            command.setBytes(pDirty, atIndex: 5)
+            command.setBuffer(layerPrev.delta.metal, atIndex: 6)
+
+            command.dispatchThreads(
+                width: width * nbChannels,
+                height: height * batchSize
+            )
+            command.enqueue()
+
+            propagateDirty()
+        }
+    }
+}

Sources/MAKit/Layer2D/MaxPool2D.swift

@@ -433,6 +433,7 @@ public class MaxPool2D: Layer2D
                         deltaCur
                 }}
             }}
+            propagateDirty()
         }
     }