Add a input-free version of pre- and post-hooks that can manipulate the module itself.

Author: NiklasGustafsson

src/TorchSharp/NN/Module.cs

@@ -1216,7 +1216,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) { }
 
@@ -1237,22 +1237,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;
@@ -1263,7 +1282,7 @@ public void remove()
                         module.remove(key);
                     }
 
-                    private HookableModule<TPreHook, TPostHook> module;
+                    private HookableModule<TModule,TPreHook, TPostHook> module;
                     private string key;
                 }
             }
@@ -1273,7 +1292,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) { }
@@ -1293,6 +1312,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)
@@ -1309,6 +1332,10 @@ public TResult call(T input)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1319,7 +1346,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) { }
@@ -1339,6 +1366,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) {
@@ -1357,6 +1388,10 @@ public TResult call(T1 input1, T2 input2)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1368,7 +1403,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) { }
@@ -1388,6 +1423,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) {
@@ -1407,6 +1446,10 @@ public TResult call(T1 input1, T2 input2, T3 input3)
                             result = modified;
                     }
 
+                    foreach (var hook in module_post_hooks.Values) {
+                        hook(this);
+                    }
+
                     return result;
                 }
             }
@@ -1419,7 +1462,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) { }
@@ -1439,6 +1482,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) {
@@ -1459,6 +1506,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;
                 }
             }
@@ -1472,7 +1523,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) { }
@@ -1492,6 +1543,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) {
@@ -1513,6 +1568,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;
                 }
             }
@@ -1527,7 +1586,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) { }
@@ -1547,6 +1606,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) {
@@ -1569,6 +1632,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);
+        }
     }
 }