simplify NN modules to avoid storing information in .NET

Author: dsyme

src/Examples/AlexNet.cs

@@ -49,9 +49,9 @@ static void Main(string[] args)
 
         private class Model : NN.Module
         {
-            private readonly NN.Module features;
-            private readonly NN.Module avgPool;
-            private readonly NN.Module classifier;
+            private readonly NN.Sequential features;
+            private readonly NN.AdaptiveAvgPool2D avgPool;
+            private readonly NN.Sequential classifier;
 
             public Model(int numClasses)
             {
@@ -70,23 +70,24 @@ public Model(int numClasses)
                     Relu(inPlace: true),
                     MaxPool2D(kernelSize: new long[] { 2 }));
 
-                avgPool = AdaptiveAvgPool2D(2, 2);
+                avgPool = AdaptiveAvgPool2D(new long[] { 2, 2 });
 
                 classifier = Sequential(
-                    Dropout(IsTraining()),
+                    Dropout(),
                     Linear(256 * 2 * 2, 4096),
                     Relu(inPlace: true),
-                    Dropout(IsTraining()),
+                    Dropout(),
                     Linear(4096, 4096),
                     Relu(inPlace: true),
                     Linear(4096, numClasses)
                 );
 
-                RegisterModule(features);
+                RegisterModule (features);
+                RegisterModule (avgPool);
                 RegisterModule(classifier);
             }
 
-            public override TorchTensor Forward(TorchTensor input)
+            public TorchTensor Forward(TorchTensor input)
             {
                 using (var f = features.Forward(input))
                 using (var avg = avgPool.Forward(f))
@@ -97,7 +98,7 @@ public override TorchTensor Forward(TorchTensor input)
         }
 
         private static void Train(
-        NN.Module model,
+        Model model,
         NN.Optimizer optimizer,
         Loss loss,
         IEnumerable<(TorchTensor, TorchTensor)> dataLoader,
@@ -140,7 +141,7 @@ private static void Train(
         }
 
         private static void Test(
-            NN.Module model,
+            Model model,
             Loss loss,
             IEnumerable<(TorchTensor, TorchTensor)> dataLoader,
             long size)

src/Examples/MNIST.cs

@@ -45,10 +45,10 @@ static void Main(string[] args)
 
         private class Model : NN.Module
         {
-            private NN.Module conv1 = Conv2D(1, 10, 5);
-            private NN.Module conv2 = Conv2D(10, 20, 5);
-            private NN.Module fc1 = Linear(320, 50);
-            private NN.Module fc2 = Linear(50, 10);
+            private NN.Conv2D conv1 = Conv2D(1, 10, 5);
+            private NN.Conv2D conv2 = Conv2D(10, 20, 5);
+            private NN.Linear fc1 = Linear(320, 50);
+            private NN.Linear fc2 = Linear(50, 10);
 
             public Model()
             {
@@ -58,22 +58,22 @@ public Model()
                 RegisterModule(fc2);
             }
 
-            public override TorchTensor Forward(TorchTensor input)
+            public TorchTensor Forward(TorchTensor input)
             {
                 using (var l11 = conv1.Forward(input))
-                using (var l12 = MaxPool2D(l11, kernelSize: new long[]{ 2 }))
+                using (var l12 = MaxPool2D (l11, kernelSize: new long[]{ 2 }))
                 using (var l13 = Relu(l12))
 
                 using (var l21 = conv2.Forward(l13))
-                using (var l22 = FeatureDropout(l21))
-                using (var l23 = MaxPool2D(l22, kernelSize: new long[] { 2 }))
+                using (var l22 = FeatureAlphaDropout(l21))
+                using (var l23 = MaxPool2D (l22, kernelSize: new long[] { 2 }))
                 using (var l24 = Relu(l23))
 
                 using (var x = l24.View(new long[] { -1, 320 }))
 
                 using (var l31 = fc1.Forward(x))
                 using (var l32 = Relu(l31))
-                using (var l33 = Dropout(l32, IsTraining()))
+                using (var l33 = Dropout(l32))
 
                 using (var l41 = fc2.Forward(l33))
 
@@ -82,7 +82,7 @@ public override TorchTensor Forward(TorchTensor input)
         }
 
         private static void Train(
-            NN.Module model,
+            Model model,
             NN.Optimizer optimizer,
             Loss loss,
             IEnumerable<(TorchTensor, TorchTensor)> dataLoader,
@@ -119,7 +119,7 @@ private static void Train(
         }
 
         private static void Test(
-            NN.Module model,
+            Model model,
             Loss loss,
             IEnumerable<(TorchTensor, TorchTensor)> dataLoader,
             long size)