✨ feat(core): initForward,Backward model API (#109)

owkin · Dec 2, 2023 · 516833d · 516833d
1 parent 3130f05
commit 516833d
Show file tree

Hide file tree

Showing 8 changed files with 120 additions and 41 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,7 +4,8 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
-🪜 **feat:** Dropout1D ([#108](https://github.com/owkin/GrAIdient/pull/108))\
+⚙️ **core:** initForward,Backward model API ([109](https://github.com/owkin/GrAIdient/pull/109))\
+🪜 **layer_1d:** Dropout1D ([#108](https://github.com/owkin/GrAIdient/pull/108))\
 🪜 **feat:** VQGrad, VQGradSeq ([#107](https://github.com/owkin/GrAIdient/pull/107))
 
 ## 0.3.1 (2023-08-09)

diff --git a/Sources/GrAIdient/Core/Layer/Layer.swift b/Sources/GrAIdient/Core/Layer/Layer.swift
@@ -271,6 +271,27 @@ open class Layer: Codable
     ///
     open func initKernelGPU() {}
 
+    ///
+    /// Initialize state resources in the CPU execution context.
+    ///
+    /// We initialize the neurons' state (forward and backward).
+    ///
+    open func checkStateCPU(batchSize: Int) throws {}
+
+    ///
+    /// Initialize state resources in the GPU execution context.
+    ///
+    /// We initialize the neurons' forward state.
+    ///
+    open func checkStateForwardGPU(batchSize: Int) throws {}
+
+    ///
+    /// Initialize state resources in the GPU execution context.
+    ///
+    /// We initialize the neurons' backward state.
+    ///
+    open func checkStateBackwardGPU(batchSize: Int) throws {}
+
     ///
     /// Update the backward dirty flag for `layerPrev` instance.
     ///

diff --git a/Sources/GrAIdient/Core/Model/Model.swift b/Sources/GrAIdient/Core/Model/Model.swift
@@ -682,6 +682,45 @@ public class Model: BaseModel
         }
     }
 
+    ///
+    /// Initialize state resources.
+    ///
+    /// We initialize the neurons' forward's state.
+    ///
+    public func initForward(batchSize: Int) throws
+    {
+        if GrAI.Opti.GPU
+        {
+            for layer in layers
+            {
+                try layer.checkStateForwardGPU(batchSize: batchSize)
+            }
+        }
+        else
+        {
+            for layer in layers
+            {
+                try layer.checkStateCPU(batchSize: batchSize)
+            }
+        }
+    }
+
+    ///
+    /// Initialize state resources.
+    ///
+    /// We initialize the neurons' backward's state.
+    ///
+    public func initBackward(batchSize: Int) throws
+    {
+        if GrAI.Opti.GPU
+        {
+            for layer in layers
+            {
+                try layer.checkStateBackwardGPU(batchSize: batchSize)
+            }
+        }
+    }
+
     ///
     /// Initialize hard resources and set the parameters for the optimizer.
     ///

diff --git a/Sources/GrAIdient/Layer1D/Base/Layer1D.swift b/Sources/GrAIdient/Layer1D/Base/Layer1D.swift
@@ -113,7 +113,7 @@ open class Layer1D: Layer
     ///
     /// We initialize the neurons' state (forward and backward).
     ///
-    public func checkStateCPU(batchSize: Int) throws
+    public override func checkStateCPU(batchSize: Int) throws
     {
         if neurons.nbElems == 0
         {
@@ -134,7 +134,7 @@ open class Layer1D: Layer
     ///
     /// We initialize the neurons' forward state.
     ///
-    public func checkStateForwardGPU(batchSize: Int) throws
+    public override func checkStateForwardGPU(batchSize: Int) throws
     {
         if outs == nil
         {
@@ -153,17 +153,20 @@ open class Layer1D: Layer
     ///
     /// We initialize the neurons' backward state.
     ///
-    public func checkStateBackwardGPU(batchSize: Int) throws
+    public override func checkStateBackwardGPU(batchSize: Int) throws
     {
-        if delta == nil
+        if computeDelta
         {
-            delta = MetalPrivateBuffer<Float>(
-                batchSize * nbNeurons, deviceID: deviceID
-            )
-        }
-        else if batchSize <= 0 || batchSize > delta.nbElems / nbNeurons
-        {
-            throw LayerError.BatchSize
+            if delta == nil
+            {
+                delta = MetalPrivateBuffer<Float>(
+                    batchSize * nbNeurons, deviceID: deviceID
+                )
+            }
+            else if batchSize <= 0 || batchSize > delta.nbElems / nbNeurons
+            {
+                throw LayerError.BatchSize
+            }
         }
     }
 

diff --git a/Sources/GrAIdient/Layer2D/Base/Layer2D.swift b/Sources/GrAIdient/Layer2D/Base/Layer2D.swift
@@ -162,7 +162,7 @@ open class Layer2D: Layer
     ///
     /// We initialize the neurons' state (forward and backward).
     ///
-    public func checkStateCPU(batchSize: Int) throws
+    public override func checkStateCPU(batchSize: Int) throws
     {
         if neurons.count == 0
         {
@@ -188,7 +188,7 @@ open class Layer2D: Layer
     ///
     /// We initialize the neurons' forward state.
     ///
-    public func checkStateForwardGPU(batchSize: Int) throws
+    public override func checkStateForwardGPU(batchSize: Int) throws
     {
         if outs == nil
         {
@@ -208,18 +208,21 @@ open class Layer2D: Layer
     ///
     /// We initialize the neurons' backward state.
     ///
-    public func checkStateBackwardGPU(batchSize: Int) throws
+    public override func checkStateBackwardGPU(batchSize: Int) throws
     {
-        if delta == nil
+        if computeDelta
         {
-            delta = MetalPrivateBuffer<Float>(
-                batchSize * nbChannels * width * height, deviceID: deviceID
-            )
-        }
-        else if batchSize <= 0 ||
-                batchSize > delta.nbElems / (nbChannels * width * height)
-        {
-            throw LayerError.BatchSize
+            if delta == nil
+            {
+                delta = MetalPrivateBuffer<Float>(
+                    batchSize * nbChannels * width * height, deviceID: deviceID
+                )
+            }
+            else if batchSize <= 0 ||
+                    batchSize > delta.nbElems / (nbChannels * width * height)
+            {
+                throw LayerError.BatchSize
+            }
         }
     }
 

diff --git a/Sources/GrAIdient/Layer2D/Convolution2D.swift b/Sources/GrAIdient/Layer2D/Convolution2D.swift
@@ -791,6 +791,12 @@ public class Convolution2D: BN2D, LayerWeightInit
         let weightsPtr = _wBuffers.w_p!.shared.buffer
         let biasesPtr = _bBuffers.w_p!.shared.buffer
 
+        /*let data = Data(
+            bytes: _weightsList,
+            count: nbWeights*weightHeight*weightWidth*MemoryLayout<Float>.size
+        )
+        _ = data.copyBytes(to: weightsPtr)*/
+
         for elem in 0..<nbWeights * weightHeight * weightWidth
         {
             weightsPtr[elem] = _weightsList[elem]

diff --git a/Sources/GrAIdient/Layer2D/Normalize2D.swift b/Sources/GrAIdient/Layer2D/Normalize2D.swift
@@ -418,13 +418,16 @@ public class Normalize122D: Layer2D
     ///
     public override func checkStateBackwardGPU(batchSize: Int) throws
     {
-        if _deltaTmp == nil
+        if computeDelta
         {
-            _deltaTmp = MetalPrivateBuffer<Float>(
-                batchSize * nbThreadgroups, deviceID: deviceID
-            )
+            if _deltaTmp == nil
+            {
+                _deltaTmp = MetalPrivateBuffer<Float>(
+                    batchSize * nbThreadgroups, deviceID: deviceID
+                )
+            }
+            try super.checkStateBackwardGPU(batchSize: batchSize)
         }
-        try super.checkStateBackwardGPU(batchSize: batchSize)
     }
 
     ///

diff --git a/Sources/GrAIdient/LayerSeq/Base/LayerSeq.swift b/Sources/GrAIdient/LayerSeq/Base/LayerSeq.swift
@@ -123,7 +123,7 @@ open class LayerSeq: Layer
     ///
     /// We initialize the neurons' state (forward and backward).
     ///
-    public func checkStateCPU(batchSize: Int) throws
+    public override func checkStateCPU(batchSize: Int) throws
     {
         if neurons == nil
         {
@@ -144,7 +144,7 @@ open class LayerSeq: Layer
     ///
     /// We initialize the neurons' forward state.
     ///
-    public func checkStateForwardGPU(batchSize: Int) throws
+    public override func checkStateForwardGPU(batchSize: Int) throws
     {
         if outs == nil
         {
@@ -163,18 +163,21 @@ open class LayerSeq: Layer
     ///
     /// We initialize the neurons' backward state.
     ///
-    public func checkStateBackwardGPU(batchSize: Int) throws
+    public override func checkStateBackwardGPU(batchSize: Int) throws
     {
-        if delta == nil
+        if computeDelta
         {
-            delta = MetalPrivateBuffer<Float>(
-                batchSize * sequence * nbNeurons, deviceID: deviceID
-            )
-        }
-        else if batchSize <= 0 ||
-                batchSize > delta.nbElems / (sequence * nbNeurons)
-        {
-            throw LayerError.BatchSize
+            if delta == nil
+            {
+                delta = MetalPrivateBuffer<Float>(
+                    batchSize * sequence * nbNeurons, deviceID: deviceID
+                )
+            }
+            else if batchSize <= 0 ||
+                        batchSize > delta.nbElems / (sequence * nbNeurons)
+            {
+                throw LayerError.BatchSize
+            }
         }
     }
 }