CuDnnPoolingLayer backpropagation switched to cuDNN

Author: Sergio0694

NeuralNetwork.NET.Cuda/Layers/CuDnnPoolingLayer.cs

@@ -1,4 +1,5 @@
-using Alea;
+using System;
+using Alea;
 using Alea.cuDNN;
 using NeuralNetworkNET.Extensions;
 using NeuralNetworkNET.Cuda.Extensions;
@@ -17,7 +18,7 @@ namespace NeuralNetworkNET.Cuda.Layers
     /// A pooling layer running on cuDNN, with a custom pooling mode
     /// </summary>
     [JsonObject(MemberSerialization.OptIn)]
-    internal sealed class CuDnnPoolingLayer : PoolingLayer
+    internal sealed class CuDnnPoolingLayer : PoolingLayer, IDisposable
     {
         #region cuDNN fields
 
@@ -41,6 +42,16 @@ internal sealed class CuDnnPoolingLayer : PoolingLayer
 
         #endregion
 
+        #region Fields
+
+        // A copy of the layer inputs
+        private Tensor _X;
+
+        // A copy of the layer output activity
+        private Tensor _Z;
+
+        #endregion
+
         public CuDnnPoolingLayer(in TensorInfo input, in PoolingInfo operation, ActivationFunctionType activation) : base(input, operation, activation)
         {
             PoolingDescription.Set2D((PoolingMode)operation.Mode, NanPropagation.PROPAGATE_NAN, operation.WindowHeight, operation.WindowWidth, operation.VerticalPadding, operation.HorizontalPadding, operation.VerticalStride, operation.HorizontalStride);
@@ -49,6 +60,8 @@ public CuDnnPoolingLayer(in TensorInfo input, in PoolingInfo operation, Activati
         /// <inheritdoc/>
         public override void Forward(in Tensor x, out Tensor z, out Tensor a)
         {
+            _X.TryFree();
+            x.Duplicate(out _X);
             using (DeviceMemory<float>
                 x_gpu = DnnInstance.Gpu.AllocateDevice(x),
                 z_gpu = DnnInstance.Gpu.AllocateDevice<float>(x.Entities * OutputInfo.Size))
@@ -58,6 +71,8 @@ public override void Forward(in Tensor x, out Tensor z, out Tensor a)
                 OutputDescription.Set4D(DataType.FLOAT, TensorFormat.CUDNN_TENSOR_NCHW, x.Entities, OutputInfo.Channels, OutputInfo.Height, OutputInfo.Width);
                 DnnInstance.PoolingForward(PoolingDescription, 1, InputDescription, x_gpu.Ptr, 0, OutputDescription, z_gpu.Ptr);
                 z_gpu.CopyToHost(x.Entities, OutputInfo.Size, out z);
+                _Z.TryFree();
+                z.Duplicate(out _Z);
 
                 // Activation
                 DnnInstance.ActivationForward(z.Entities, z.Length, z_gpu.Ptr, z_gpu.Ptr, ActivationFunctions.Activation);
@@ -66,7 +81,24 @@ public override void Forward(in Tensor x, out Tensor z, out Tensor a)
         }
 
         /// <inheritdoc/>
-        public override void Backpropagate(in Tensor delta_1, in Tensor z, ActivationFunction activationPrime) => z.UpscalePool2x2(delta_1, InputInfo.Channels);
+        public override void Backpropagate(in Tensor delta_1, in Tensor z, ActivationFunction activationPrime)
+        {
+            using (DeviceMemory<float> dx_gpu = DnnInstance.Gpu.AllocateDevice<float>(z.Size))
+            {
+                using (DeviceMemory<float>
+                    x_gpu = DnnInstance.Gpu.AllocateDevice(_X),
+                    y_gpu = DnnInstance.Gpu.AllocateDevice(_Z),
+                    dy_gpu = DnnInstance.Gpu.AllocateDevice(delta_1))
+                {
+                    DnnInstance.PoolingBackward(PoolingDescription, 1, OutputDescription, y_gpu.Ptr, OutputDescription, dy_gpu.Ptr, InputDescription, x_gpu.Ptr, 0, InputDescription, dx_gpu.Ptr);
+                }
+                using (DeviceMemory<float> z_gpu = DnnInstance.Gpu.AllocateDevice(z))
+                {
+                    DnnInstance.ActivationBackward(z.Entities, z.Length, z_gpu.Ptr, dx_gpu.Ptr, activationPrime);
+                    z_gpu.CopyTo(z);
+                }
+            }
+        }
 
         /// <inheritdoc/>
         public override INetworkLayer Clone() => new CuDnnPoolingLayer(InputInfo, OperationInfo, ActivationFunctionType);
@@ -84,5 +116,25 @@ public override void Forward(in Tensor x, out Tensor z, out Tensor a)
             if (!stream.TryRead(out PoolingInfo operation)) return null;
             return new CuDnnPoolingLayer(input, operation, activation);
         }
+
+        #region IDisposable
+
+        ~CuDnnPoolingLayer() => Dispose();
+
+        /// <inheritdoc/>
+        void IDisposable.Dispose()
+        {
+            GC.SuppressFinalize(this);
+            Dispose();
+        }
+
+        // Private Dispose method
+        private void Dispose()
+        {
+            _X.TryFree();
+            _Z.TryFree();
+        }
+
+        #endregion
     }
 }

Unit/NeuralNetwork.NET.Cuda.Unit/CuDnnLayersTest.cs

@@ -236,15 +236,41 @@ public void PoolingForward()
         }
 
         [TestMethod]
-        public void PoolingBackward()
+        public unsafe void PoolingBackward()
         {
-            float[,]
-                delta_1 = WeightsProvider.NewFullyConnectedWeights(TensorInfo.CreateLinear(400), 29 * 29 * 3, WeightsInitializationMode.GlorotNormal).AsMatrix(400, 29 * 29 * 3),
-                z = WeightsProvider.NewFullyConnectedWeights(TensorInfo.CreateLinear(400), 58 * 58 * 3, WeightsInitializationMode.GlorotNormal).AsMatrix(400, 58 * 58 * 3);
+            // Setup
+            Tensor.New(400, 58 * 58 * 3, out Tensor x);
+            KerasWeightsProvider.FillWithHeEtAlUniform(x, 10);
             PoolingLayer
                 cpu = new PoolingLayer(new TensorInfo(58, 58, 3), PoolingInfo.Default, ActivationFunctionType.LeakyReLU),
                 gpu = new CuDnnPoolingLayer(cpu.InputInfo, PoolingInfo.Default, ActivationFunctionType.LeakyReLU);
-            TestBackward(cpu, gpu, delta_1, z);
+            gpu.Forward(x, out Tensor z, out Tensor a);
+            a.Free();
+            x.Duplicate(out Tensor x2);
+            Tensor.New(z.Entities, z.Length, out Tensor delta);
+            KerasWeightsProvider.FillWithHeEtAlUniform(delta, 10);
+
+            // Backward
+            cpu.Backpropagate(delta, x, ActivationFunctions.LeakyReLUPrime);
+            gpu.Backpropagate(delta, x2, ActivationFunctions.LeakyReLUPrime);
+            bool valid = true;
+            float* px = (float*)x.Ptr.ToPointer(), px2 = (float*)x2.Ptr.ToPointer();
+            int count = 0;
+            for (int i = 0; i < x.Size; i++)
+            {
+                if (px[i].EqualsWithDelta(px2[i], 1e-5f)) continue;
+                if (px[i].EqualsWithDelta(px2[i] * 100f, 1e-5f)) count++;   // The cuDNN pooling backwards method returns a value scaled by 0.01 from time to time for some reason (less than 2% anyways)
+                else
+                {
+                    valid = false;
+                    break;
+                }
+            }
+            Assert.IsTrue(valid && count * 100f / x.Size < 2);
+            x.Free();
+            x2.Free();
+            z.Free();
+            delta.Free();
         }
 
         #endregion