Merge pull request #1318 from NiklasGustafsson/missing

Add a input-free version of pre- and post-hooks that can manipulate t…

Author: NiklasGustafsson

RELEASENOTES.md

@@ -6,32 +6,33 @@ Releases, starting with 9/2/2021, are listed with the most recent release at the
 
 __Breaking Changes__:
 
-- `torchvision.dataset.MNIST` will try more mirrors.
-    - The thrown exception might be changed when it fails to download `MNIST`, `FashionMNIST` or `KMNIST`.
-- `ObjectDisposedException` will now be thrown when trying to use the disposed dispose scopes.
-    - The constructor of dispose scopes is no longer `public`. Use `torch.NewDisposeScope` instead.
+`torchvision.dataset.MNIST` will try more mirrors. The thrown exception might be changed when it fails to download `MNIST`, `FashionMNIST` or `KMNIST`.<br/>
+`ObjectDisposedException` will now be thrown when trying to use a disposed dispose scopes.<br/>
+The constructor of dispose scopes is no longer `public`. Use `torch.NewDisposeScope` instead.<br/>
 
 __API Changes__:
 
-- #1291 `Tensor.grad()` and `Tensor.set_grad()` have been replaced by a new property `Tensor.grad`.
-    - A potential memory leak caused by `set_grad` has been resolved.
-- `Include` method of dispose scopes has been removed. Use `Attach` instead.
-- Two more `Attach` methods that accepts `IEnumerable<IDisposable>`s and arrays as the parameter have been added into dispose scopes.
-- A new property `torch.CurrentDisposeScope` has been added to provide the ability to get the current dispose scope.
+#1314 Grant read-only access to DataLoader attributes<br/>
+#1313 Add 'non_blocking' argument to tensor and module 'to()' signatures.<br/>
+#1291 `Tensor.grad()` and `Tensor.set_grad()` have been replaced by a new property `Tensor.grad`.<br/>
+A potential memory leak caused by `set_grad` has been resolved.<br/>
+`Include` method of dispose scopes has been removed. Use `Attach` instead.<br/>
+Two more `Attach` methods that accepts `IEnumerable<IDisposable>`s and arrays as the parameter have been added into dispose scopes.<br/>
+A new property `torch.CurrentDisposeScope` has been added to provide the ability to get the current dispose scope.<br/>
+Add module hooks that take no input/output arguments, just the module itself.<br/>
 
 __Bug Fixes__:
 
-- #1300 `Adadelta`, `Adam` and `AdamW` will no longer throw `NullReferenceException` when `maximize` is `true` and `grad` is `null`.
-- `torch.normal` will now correctly return a leaf tensor.
-- New options `disposeBatch` and `disposeDataset` have been added into `DataLoader`.
-    - The default collate functions will now always dispose the intermediate tensors, rather than wait for the next iteration.
-    - The fields are now exposed as readonly properties.
+#1300 `Adadelta`, `Adam` and `AdamW` will no longer throw `NullReferenceException` when `maximize` is `true` and `grad` is `null`.<br/>
+torch.normal` will now correctly return a leaf tensor.<br/>
+New options `disposeBatch` and `disposeDataset` have been added into `DataLoader`.<br/>
+The default collate functions will now always dispose the intermediate tensors, rather than wait for the next iteration.<br/>
 
 __Bug Fixes__:
 
-- `TensorDataset` will now keep the aliases detached from dispose scopes, to avoid the unexpected disposal.
-- `DataLoaderEnumerator` has been completely rewritten to resolve the unexpected shuffler disposal, the ignorance of `drop_last`, the incorrect count of worker, and the potential leak cause by multithreading.
-- #1303 Allow dispose scopes to be disposed out of LIFO order.
+`TensorDataset` will now keep the aliases detached from dispose scopes, to avoid the unexpected disposal.<br/>
+`DataLoaderEnumerator` has been completely rewritten to resolve the unexpected shuffler disposal, the ignorance of `drop_last`, the incorrect count of worker, and the potential leak cause by multithreading.<br/>
+#1303 Allow dispose scopes to be disposed out of LIFO order.<br/>
 
 # NuGet Version 0.102.4
 

src/TorchSharp/NN/Module.cs

@@ -1228,7 +1228,7 @@ public interface IModule<T1, T2, T3, T4, T5, T6, TResult>
             /// <summary>
             /// Represents a module that accepts 'hook' to the module logic.
             /// </summary>
-            public class HookableModule<TPreHook,TPostHook> : Module
+            public class HookableModule<TModule,TPreHook,TPostHook> : Module
             {
                 protected HookableModule(string name) : base(name) { }
 
@@ -1249,22 +1249,41 @@ public HookRemover register_forward_pre_hook(TPreHook hook)
                     return new HookRemover(this, key);
                 }
 
+                public HookRemover register_forward_hook(Action<TModule> hook)
+                {
+                    var key = Guid.NewGuid().ToString();
+                    module_post_hooks.Add(key, hook);
+                    return new HookRemover(this, key);
+                }
+
+                public HookRemover register_forward_pre_hook(Action<TModule> hook)
+                {
+                    var key = Guid.NewGuid().ToString();
+                    module_pre_hooks.Add(key, hook);
+                    return new HookRemover(this, key);
+                }
+
                 private void remove(string key)
                 {
                     if (pre_hooks.ContainsKey(key)) pre_hooks.Remove(key);
                     if (post_hooks.ContainsKey(key)) post_hooks.Remove(key);
+                    if (module_pre_hooks.ContainsKey(key)) module_pre_hooks.Remove(key);
+                    if (module_post_hooks.ContainsKey(key)) module_post_hooks.Remove(key);
                 }
 
                 protected Dictionary<string, TPreHook> pre_hooks = new Dictionary<string, TPreHook>();
                 protected Dictionary<string, TPostHook> post_hooks = new Dictionary<string, TPostHook>();
 
+                protected Dictionary<string, Action<TModule>> module_pre_hooks = new Dictionary<string, Action<TModule>>();
+                protected Dictionary<string, Action<TModule>> module_post_hooks = new Dictionary<string, Action<TModule>>();
+
                 /// <summary>
                 /// Used to remove a specific hook, following the PyTorch API design.
                 /// </summary>
                 /// <remarks>The name and namespace of this class is not the same as in PyTorch, but serves the same purpose.</remarks>
                 public class HookRemover
                 {
-                    public HookRemover(HookableModule<TPreHook, TPostHook> module, string key)
+                    public HookRemover(HookableModule<TModule,TPreHook, TPostHook> module, string key)
                     {
                         this.module = module;
                         this.key = key;
@@ -1275,7 +1294,7 @@ public void remove()
                         module.remove(key);
                     }
 
-                    private HookableModule<TPreHook, TPostHook> module;
+                    private HookableModule<TModule,TPreHook, TPostHook> module;
                     private string key;
                 }
             }
@@ -1285,7 +1304,7 @@ public void remove()
             /// </summary>
             /// <typeparam name="T">The argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T, TResult> : HookableModule<Func<Module<T,TResult>, T, T>, Func<Module<T, TResult>, T, TResult, TResult>>, IModule<T, TResult>
+            public abstract class Module<T, TResult> : HookableModule<Module<T, TResult>, Func<Module<T,TResult>, T, T>, Func<Module<T, TResult>, T, TResult, TResult>>, IModule<T, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1305,6 +1324,10 @@ public TResult call(T input)
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input);
                         if (modified is not null)
@@ -1321,6 +1344,10 @@ public TResult call(T input)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1331,7 +1358,7 @@ public TResult call(T input)
             /// <typeparam name="T1">The first argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T2">The second argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, TResult> : HookableModule<Func<Module<T1, T2, TResult>, T1, T2, (T1, T2)?>, Func<Module<T1, T2, TResult>, T1, T2, TResult, TResult>>, IModule<T1, T2, TResult>
+            public abstract class Module<T1, T2, TResult> : HookableModule<Module<T1, T2, TResult>, Func<Module<T1, T2, TResult>, T1, T2, (T1, T2)?>, Func<Module<T1, T2, TResult>, T1, T2, TResult, TResult>>, IModule<T1, T2, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1351,6 +1378,10 @@ public TResult call(T1 input1, T2 input2)
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input1, input2);
                         if (modified.HasValue) {
@@ -1369,6 +1400,10 @@ public TResult call(T1 input1, T2 input2)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1380,7 +1415,7 @@ public TResult call(T1 input1, T2 input2)
             /// <typeparam name="T2">The second argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T3">The third argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, TResult> : HookableModule<Func<Module<T1, T2, T3, TResult>, T1, T2, T3, (T1, T2, T3)?>, Func<Module<T1, T2, T3, TResult>, T1, T2, T3, TResult, TResult>>, IModule<T1, T2, T3, TResult>
+            public abstract class Module<T1, T2, T3, TResult> : HookableModule<Module<T1, T2, T3, TResult>,Func<Module<T1, T2, T3, TResult>, T1, T2, T3, (T1, T2, T3)?>, Func<Module<T1, T2, T3, TResult>, T1, T2, T3, TResult, TResult>>, IModule<T1, T2, T3, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1400,6 +1435,10 @@ public TResult call(T1 input1, T2 input2, T3 input3)
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input1, input2, input3);
                         if (modified.HasValue) {
@@ -1419,6 +1458,10 @@ public TResult call(T1 input1, T2 input2, T3 input3)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1431,7 +1474,7 @@ public TResult call(T1 input1, T2 input2, T3 input3)
             /// <typeparam name="T3">The third argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T4">The fourth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, TResult> : HookableModule<Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, (T1, T2, T3, T4)?>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, TResult, TResult>>, IModule<T1, T2, T3, T4, TResult>
+            public abstract class Module<T1, T2, T3, T4, TResult> : HookableModule<Module<T1, T2, T3, T4, TResult>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, (T1, T2, T3, T4)?>, Func<Module<T1, T2, T3, T4, TResult>, T1, T2, T3, T4, TResult, TResult>>, IModule<T1, T2, T3, T4, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1451,6 +1494,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4)
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input1, input2, input3, input4);
                         if (modified.HasValue) {
@@ -1471,6 +1518,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1484,7 +1535,7 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4)
             /// <typeparam name="T4">The fourth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T5">The fifth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, T5,TResult> : HookableModule<Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, (T1, T2, T3, T4, T5)?>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, TResult>
+            public abstract class Module<T1, T2, T3, T4, T5,TResult> : HookableModule<Module<T1, T2, T3, T4, T5,TResult>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, (T1, T2, T3, T4, T5)?>, Func<Module<T1, T2, T3, T4, T5, TResult>, T1, T2, T3, T4, T5, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1504,6 +1555,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5)
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input1, input2, input3, input4, input5);
                         if (modified.HasValue) {
@@ -1525,6 +1580,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1539,7 +1598,7 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5)
             /// <typeparam name="T5">The fifth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="T6">The sixth argument type of the module's forward() function.</typeparam>
             /// <typeparam name="TResult">The return type of the module's forward() function.</typeparam>
-            public abstract class Module<T1, T2, T3, T4, T5, T6, TResult> : HookableModule<Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, (T1, T2, T3, T4, T5, T6)?>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, T6, TResult>
+            public abstract class Module<T1, T2, T3, T4, T5, T6, TResult> : HookableModule<Module<T1, T2, T3, T4, T5, T6, TResult>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, (T1, T2, T3, T4, T5, T6)?>, Func<Module<T1, T2, T3, T4, T5, T6, TResult>, T1, T2, T3, T4, T5, T6, TResult, TResult>>, IModule<T1, T2, T3, T4, T5, T6, TResult>
             {
                 protected Module(string name) : base(name) { }
                 protected Module(IntPtr handle, IntPtr boxedHandle) : base(handle, boxedHandle) { }
@@ -1559,6 +1618,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5, T6 in
                 {
                     // Call pre-hooks, if available.
 
+                    foreach (var hook in module_pre_hooks.Values) {
+                        hook(this);
+                    }
+                    
                     foreach (var hook in pre_hooks.Values) {
                         var modified = hook(this, input1, input2, input3, input4, input5, input6);
                         if (modified.HasValue) {
@@ -1581,6 +1644,10 @@ public TResult call(T1 input1, T2 input2, T3 input3, T4 input4, T5 input5, T6 in
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }

test/TorchSharpTest/NN.cs

@@ -6583,5 +6583,45 @@ public void TestModulePostHooks()
             lin1.call(input);
             Assert.Equal(1, counter);
         }
+
+        [Fact]
+        public void TestModulePreHooksGeneric()
+        {
+            var lin1 = torch.nn.Linear(100, 10);
+            var input = torch.randn(32, 100, 100);
+            var counter = 0;
+
+            var pre_hook = (Module<Tensor, Tensor> m) => { counter += 1;};
+
+            var handle = lin1.register_forward_pre_hook(pre_hook);
+
+            lin1.call(input);
+            Assert.Equal(1, counter);
+
+            handle.remove();
+
+            lin1.call(input);
+            Assert.Equal(1, counter);
+        }
+
+        [Fact]
+        public void TestModulePostHooksGeneric()
+        {
+            var lin1 = torch.nn.Linear(100, 10);
+            var input = torch.randn(32, 100, 100);
+            var counter = 0;
+
+            var hook = (Module<Tensor, Tensor> m) => { counter += 1;};
+
+            var handle = lin1.register_forward_hook(hook);
+
+            lin1.call(input);
+            Assert.Equal(1, counter);
+
+            handle.remove();
+
+            lin1.call(input);
+            Assert.Equal(1, counter);
+        }
     }
 }