Merge pull request #2 from gyeongin/auto-grad-mode

Implement AutoGradMode for TorchSharp

Author: interesaaat

Test/TorchSharp.cs

@@ -337,6 +337,31 @@ public void TestGrad()
             }
         }
 
+        [TestMethod]
+        public void TestAutoGradMode()
+        {
+            var x = FloatTensor.RandomN(new long[] { 2, 3 }, device: "cpu:0", requiresGrad: true);
+            using (var mode = new AutoGradMode(false))
+            {
+                var sum = x.Sum();
+                sum.Backward();
+                var grad = x.Grad();
+                Assert.IsTrue(grad.Handle == IntPtr.Zero);
+            }
+            using (var mode = new AutoGradMode(true))
+            {
+                var sum = x.Sum();
+                sum.Backward();
+                var grad = x.Grad();
+                Assert.IsFalse(grad.Handle == IntPtr.Zero);
+                var data = grad.Data;
+                for (int i = 0; i < 2 * 3; i++)
+                {
+                    Assert.AreEqual(data[i], 1.0);
+                }
+            }
+        }
+
         [TestMethod]
         public void TestSubInPlace()
         {

TorchSharp/Torch.cs

@@ -1,4 +1,5 @@
-using System.Runtime.InteropServices;
+using System;
+using System.Runtime.InteropServices;
 
 namespace TorchSharp
 {
@@ -12,4 +13,35 @@ public static void SetSeed(long seed)
             NN_Seed(seed);
         }
     }
+
+    public class AutoGradMode : IDisposable
+    {
+        [DllImport("LibTorchSharp")]
+        extern static bool THS_gradmode_is_enabled();
+
+        [DllImport("LibTorchSharp")]
+        extern static void THS_gradmode_set_enabled(bool enabled);
+
+        public AutoGradMode(bool enabled)
+        {
+            prev_mode = THS_gradmode_is_enabled();
+            THS_gradmode_set_enabled(enabled);
+        }
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        public void Dispose(bool disposing)
+        {
+            if (disposing)
+            {
+                THS_gradmode_set_enabled(prev_mode);
+            }
+        }
+
+        bool prev_mode;
+    }
 }