Added modules for MNIST model

Author: interesaaat

TorchSharp/NN/Conv2D.cs

@@ -0,0 +1,21 @@
+using System;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    public class Conv2D : Module
+    {
+        internal Conv2D(IntPtr handle) : base(handle)
+        {
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static FloatTensor.HType NN_conv2DModule_Forward(Module.HType module, IntPtr tensor);
+
+        public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
+        {
+            return new FloatTensor(NN_conv2DModule_Forward(handle, tensor.Handle));
+        }
+    }
+}

TorchSharp/NN/Dropout.cs

@@ -0,0 +1,29 @@
+using System;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    /// <summary>
+    /// This class is used to represent a dropout module.
+    /// </summary>
+    public class Dropout : FunctionalModule
+    {
+        private double _probability;
+        private Func<bool> _isTraining;
+
+        internal Dropout(double probability, Func<bool> isTraining) : base()
+        {
+            _probability = probability;
+            _isTraining = isTraining;
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static FloatTensor.HType NN_LogSoftMaxModule_Forward(IntPtr tensor, double probability, bool isTraining);
+
+        public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
+        {
+            return new FloatTensor(NN_LogSoftMaxModule_Forward(tensor.Handle, _probability, _isTraining.Invoke()));
+        }
+    }
+}

TorchSharp/NN/Functional.cs

@@ -0,0 +1,36 @@
+using System;
+using System.Collections.Generic;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    /// <summary>
+    /// This class is used to represent a functional module (e.g., ReLU).
+    /// </summary>
+    public abstract class FunctionalModule : Module
+    {
+        internal FunctionalModule() : base(IntPtr.Zero)
+        {
+        }
+
+        public override void ZeroGrad()
+        {
+        }
+
+        public override bool IsTraining()
+        {
+            return true;
+        }
+
+        public override IEnumerable<ITorchTensor<float>> Parameters()
+        {
+            return new List<ITorchTensor<float>>();
+        }
+
+        public override string[] GetModules()
+        {
+            return new string[0];
+        }
+    }
+}

TorchSharp/NN/FunctionalModule.cs

@@ -0,0 +1,27 @@
+using System;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    /// <summary>
+    /// This class is used to represent a log softmax module.
+    /// </summary>
+    public class LogSoftMax : FunctionalModule
+    {
+        private long _dimension;
+
+        internal LogSoftMax(long dimension) : base()
+        {
+            _dimension = dimension;
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static FloatTensor.HType NN_LogSoftMaxModule_Forward(IntPtr tensor, long dimension);
+
+        public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
+        {
+            return new FloatTensor(NN_LogSoftMaxModule_Forward(tensor.Handle, _dimension));
+        }
+    }
+}

TorchSharp/NN/LogSoftMax.cs

@@ -0,0 +1,27 @@
+using System;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    /// <summary>
+    /// This class is used to represent a ReLu module.
+    /// </summary>
+    public class MaxPool2D : FunctionalModule
+    {
+        private long _kernelSize;
+
+        internal MaxPool2D(long kernelSize) : base()
+        {
+            _kernelSize = kernelSize;
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static FloatTensor.HType NN_MaxPool2DModule_Forward(IntPtr tensor, long kernelSize);
+
+        public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
+        {
+            return new FloatTensor(NN_MaxPool2DModule_Forward(tensor.Handle, _kernelSize));
+        }
+    }
+}

TorchSharp/NN/MaxPool2D.cs

@@ -73,7 +73,7 @@ protected void Dispose(bool disposing)
         }
     }
 
-    public abstract partial class Module : IDisposable
+    public partial class Module : IDisposable
     {
         static public Sequential Sequential(params Module[] modules)
         {
@@ -89,15 +89,62 @@ static public Module Linear(int input, int output, bool hasBias = false)
         }
 
         [DllImport("LibTorchSharp")]
-        extern static IntPtr NN_reluModule();
+        extern static IntPtr NN_conv2dModule(long inputChannel, long outputChannel, int kernelSize);
+
+        static public Module Conv2D(long inputChannel, long outputChannel, int kernelSize)
+        {
+            return new Conv2D(NN_conv2dModule(inputChannel, outputChannel, kernelSize));
+        }
 
         static public Module Relu()
         {
-            return new Functional(NN_reluModule());
+            return new ReLu();
+        }
+
+        static public ITorchTensor<float> Relu(ITorchTensor<float> x)
+        {
+            return new ReLu().Forward(x);
+        }
+
+        static public Module MaxPool2D(long kernelSize)
+        {
+            return new MaxPool2D(kernelSize);
+        }
+
+        static public ITorchTensor<float> MaxPool2D(ITorchTensor<float> x, long kernelSize)
+        {
+            return new MaxPool2D(kernelSize).Forward(x);
+        }
+
+        static public Module LogSoftMax(long dimension)
+        {
+            return new LogSoftMax(dimension);
+        }
+
+        static public ITorchTensor<float> LogSoftMax(ITorchTensor<float> x, long dimension)
+        {
+            return new LogSoftMax(dimension).Forward(x);
         }
 
+        static public Module Dropout(double probability, Func<bool> isTraining)
+        {
+            return new Dropout(probability, isTraining);
+        }
+
+        static public ITorchTensor<float> Dropout(ITorchTensor<float> x, double probability, Func<bool> isTraining)
+        {
+            return new Dropout(probability, isTraining).Forward(x);
+        }
+    }
+
+    public abstract partial class Module : IDisposable
+    {
         public abstract ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor);
 
+        public virtual void RegisterModule(Module module)
+        {
+        }
+
         [DllImport("LibTorchSharp")]
         extern static void NN_Module_ZeroGrad(HType module);
 
@@ -106,6 +153,14 @@ public virtual void ZeroGrad()
             NN_Module_ZeroGrad(handle);
         }
 
+        [DllImport("LibTorchSharp")]
+        extern static bool NN_IsTraining(HType module);
+
+        public virtual bool IsTraining()
+        {
+            return NN_IsTraining(handle);
+        }
+
         [DllImport("LibTorchSharp")]
         extern static void NN_GetParameters(HType module, AllocatePinnedArray allocator);
 

TorchSharp/NN/Module.cs

@@ -0,0 +1,24 @@
+using System;
+using System.Runtime.InteropServices;
+using TorchSharp.Tensor;
+
+namespace TorchSharp.NN
+{
+    /// <summary>
+    /// This class is used to represent a ReLu module.
+    /// </summary>
+    public class ReLu : FunctionalModule
+    {
+        internal ReLu() : base()
+        {
+        }
+
+        [DllImport("LibTorchSharp")]
+        extern static FloatTensor.HType NN_ReluModule_Forward(IntPtr tensor);
+
+        public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
+        {
+            return new FloatTensor(NN_ReluModule_Forward(tensor.Handle));
+        }
+    }
+}

TorchSharp/NN/ReLu.cs

@@ -18,10 +18,15 @@ public Sequential(params Module[] modules) : base(IntPtr.Zero)
         {
             foreach (var module in modules)
             {
-                Modules.Add(module);
+                RegisterModule(module);
             }
         }
 
+        public override void RegisterModule(Module module)
+        {
+            Modules.Add(module);
+        }
+
         public override ITorchTensor<float> Forward<T>(ITorchTensor<T> tensor)
         {
             if (Modules.Count < 1)