Merge pull request #1110 from NiklasGustafsson/bugs

Re-introduce `forward` to ScriptModule

Author: NiklasGustafsson

RELEASENOTES.md

@@ -6,6 +6,7 @@ Releases, starting with 9/2/2021, are listed with the most recent release at the
 
 __Bug Fixes__:
 
+ScriptModule: adding `forward` and the ability to hook.<br/>
 Update to SkiaSharp 2.88.6 to avoid the libwebp vulnerability.<br/>
 #1105: Dataset files get written to the wrong directory<br/>
 

src/TorchSharp/JIT/ScriptModule.cs

@@ -220,8 +220,8 @@ private Type GetType(Type type)
                 /// <summary>
                 /// Invoke the 'forward' function of the script with any number of arguments.
                 /// </summary>
-                /// <param name="objs"></param>
-                /// <returns></returns>
+                /// <param name="input">Any number of parameters for the forward function.</param>
+                /// <returns>An object.</returns>
                 /// <remarks>
                 /// Only certain types can currently be passed:
                 /// 1. Tensor
@@ -238,15 +238,14 @@ private Type GetType(Type type)
                 /// For returned types, if the number of values returned in a tuple is greaterh than 5, it is returned as an array, instead.
                 /// If a tuple contains both tensors and scalars, it is returned as an object[].
                 /// </remarks>
-                /// <exception cref="NotImplementedException"></exception>
-                public object call(params object[] objs)
+                public object forward(params object[] input)
                 {
                     TensorOrScalar[] ptrArray = null;
                     sbyte typeCode = 0;
 
                     using (var parray = new IndexedPinnedArrays<TensorOrScalar>()) {
 
-                        var tRefsHandle = DetermineArgumentTypeRefs(objs, out var count, parray);
+                        var tRefsHandle = DetermineArgumentTypeRefs(input, out var count, parray);
 
                         var allocated = parray.Count;
 
@@ -258,6 +257,40 @@ public object call(params object[] objs)
                     }
                 }
 
+                /// <summary>
+                /// Synonym for `forward`
+                /// </summary>
+                /// <param name="input">Any number of parameters for the forward function.</param>
+                /// <returns>An object.</returns>
+                /// <remarks>
+                /// Only certain types can currently be passed:
+                /// 1. Tensor
+                /// 2. Scalar
+                /// 3. int/long
+                /// 4. double/float
+                /// 5. bool
+                ///
+                /// Only certain types can currently be returned:
+                /// 1. Tensor / Scalar
+                /// 2. Tuple of Tensor / Scalar
+                /// 3. Array (Python list) of Tensor / Scalar
+                ///
+                /// For returned types, if the number of values returned in a tuple is greaterh than 5, it is returned as an array, instead.
+                /// If a tuple contains both tensors and scalars, it is returned as an object[].
+                ///
+                /// Note: this currently does not support hooking the module.
+                /// </remarks>
+                public object call(params object[] input)
+                {
+                    // TODO: Call pre-hooks, if available.
+
+                    var result = forward(input);
+
+                    // TODO: Call post-hooks, if available.
+
+                    return result;
+                }
+
                 /// <summary>
                 /// Invoke a function from the script module.
                 /// </summary>
@@ -437,7 +470,7 @@ internal static object ProcessReturnValue(string name, IndexedPinnedArrays<Tenso
                             // List of scalars and tensors
                             var result = new object[ptrArray.Length];
                             for (var i = 0; i < ptrArray.Length; i++) {
-                                switch(ptrArray[i].TypeCode) {
+                                switch (ptrArray[i].TypeCode) {
                                 case 0:
                                     result[i] = new Tensor(ptrArray[i].Handle);
                                     break;
@@ -566,6 +599,62 @@ internal static object ProcessReturnValue(string name, IndexedPinnedArrays<Tenso
                 public TResult invoke<T, TResult>(string name, params T[] inputs) => (TResult)invoke(name, inputs);
             }
 
+
+
+            /// <summary>
+            /// Represents a module that accepts 'hook' to the module logic.
+            /// </summary>
+            public class HookableScriptModule<TPreHook, TPostHook> : ScriptModule
+            {
+                internal HookableScriptModule(IntPtr handle) : base(handle)
+                {
+                }
+
+                public HookRemover register_forward_hook(TPostHook hook)
+                {
+                    var key = Guid.NewGuid().ToString();
+                    post_hooks.Add(key, hook);
+                    return new HookRemover(this, key);
+                }
+
+                public HookRemover register_forward_pre_hook(TPreHook hook)
+                {
+                    var key = Guid.NewGuid().ToString();
+                    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);
+                }
+
+                protected Dictionary<string, TPreHook> pre_hooks = new Dictionary<string, TPreHook>();
+                protected Dictionary<string, TPostHook> post_hooks = new Dictionary<string, TPostHook>();
+
+                /// <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(HookableScriptModule<TPreHook, TPostHook> module, string key)
+                    {
+                        this.module = module;
+                        this.key = key;
+                    }
+
+                    public void remove()
+                    {
+                        module.remove(key);
+                    }
+
+                    private HookableScriptModule<TPreHook, TPostHook> module;
+                    private string key;
+                }
+            }
+
             /// <summary>
             /// A script module taking any number of tensors as input
             /// </summary>
@@ -593,18 +682,26 @@ internal ScriptModule(IntPtr handle) : base(handle) { }
                 /// For returned types, if the number of values returned in a tuple is greaterh than 5, it is returned as an array, instead.
                 /// If a tuple contains both tensors and scalars, it is returned as an object[].
                 /// </remarks>
-                public TResult call(params Tensor[] tensor)
+                public TResult call(params Tensor[] input)
                 {
-                    return (TResult)base.call(tensor);
+                    // TODO: Call pre-hooks, if available.
+
+                    var result = forward(input);
+
+                    // TODO: Call post-hooks, if available.
+
+                    return result;
                 }
+
+                public TResult forward(params Tensor[] tensor) => (TResult)base.forward(tensor);
             }
 
             /// <summary>
             /// A script module taking a single argument.
             /// </summary>
             /// <typeparam name="T">The argument type.</typeparam>
             /// <typeparam name="TResult">The return type of the module.</typeparam>
-            public class ScriptModule<T, TResult> : ScriptModule, torch.nn.IModule<T, TResult>
+            public class ScriptModule<T, TResult> : HookableScriptModule<Func<ScriptModule<T, TResult>, T, T>, Func<ScriptModule<T, TResult>, T, TResult, TResult>>, torch.nn.IModule<T, TResult>
             {
                 internal ScriptModule(IntPtr handle) : base(handle) { }
 
@@ -627,10 +724,30 @@ internal ScriptModule(IntPtr handle) : base(handle) { }
                 /// For returned types, if the number of values returned in a tuple is greaterh than 5, it is returned as an array, instead.
                 /// If a tuple contains both tensors and scalars, it is returned as an object[].
                 /// </remarks>
-                public TResult call(T tensor)
+                public TResult call(T input)
                 {
-                    return (TResult)base.call(tensor);
+                    // Call pre-hooks, if available.
+
+                    foreach (var hook in pre_hooks.Values) {
+                        var modified = hook(this, input);
+                        if (modified is not null)
+                            input = modified;
+                    }
+
+                    var result = forward(input);
+
+                    // Call post-hooks, if available.
+
+                    foreach (var hook in post_hooks.Values) {
+                        var modified = hook(this, input, result);
+                        if (modified is not null)
+                            result = modified;
+                    }
+
+                    return result;
                 }
+
+                public TResult forward(T tensor) => (TResult)base.forward(tensor);
             }
 
             /// <summary>
@@ -639,7 +756,7 @@ public TResult call(T tensor)
             /// <typeparam name="T1">The first argument type.</typeparam>
             /// <typeparam name="T2">The second argument type.</typeparam>
             /// <typeparam name="TResult">The return type of the module.</typeparam>
-            public class ScriptModule<T1, T2, TResult> : ScriptModule, torch.nn.IModule<T1, T2, TResult>
+            public class ScriptModule<T1, T2, TResult> : HookableScriptModule<Func<ScriptModule<T1, T2, TResult>, T1, T2, (T1, T2)?>, Func<ScriptModule<T1, T2, TResult>, T1, T2, TResult, TResult>>, torch.nn.IModule<T1, T2, TResult>
             {
                 internal ScriptModule(IntPtr handle) : base(handle) { }
 
@@ -664,8 +781,30 @@ internal ScriptModule(IntPtr handle) : base(handle) { }
                 /// </remarks>
                 public TResult call(T1 input1, T2 input2)
                 {
-                    return (TResult)base.call(input1, input2);
+                    // Call pre-hooks, if available.
+
+                    foreach (var hook in pre_hooks.Values) {
+                        var modified = hook(this, input1, input2);
+                        if (modified.HasValue) {
+                            input1 = modified.Value.Item1;
+                            input2 = modified.Value.Item2;
+                        }
+                    }
+
+                    var result = forward(input1, input2);
+
+                    // Call post-hooks, if available.
+
+                    foreach (var hook in post_hooks.Values) {
+                        var modified = hook(this, input1, input2, result);
+                        if (modified is not null)
+                            result = modified;
+                    }
+
+                    return result;
                 }
+
+                public TResult forward(T1 input1, T2 input2) => (TResult)base.forward(input1, input2);
             }
 
             /// <summary>

test/TorchSharpTest/TestJIT.cs

@@ -14,6 +14,27 @@ namespace TorchSharp
     public class TestJIT
     {
 #if false
+        [Fact]
+        public void TestLoadJIT_1()
+        {
+            var input = torch.ones(10);
+            var expected = torch.tensor(new float[] { 0.313458264f, 0, 0.9996568f, 0, 0, 0 });
+
+            // One linear layer followed by ReLU.
+            var m = torch.jit.load<Tensor, Tensor>(@"linrelu.script.dat");
+            if (torch.cuda.is_available()) {
+                m = m.to(torch.CUDA);
+                input = input.to(torch.CUDA);
+                expected = expected.to(torch.CUDA);
+            }
+
+            var t = m.forward(input);
+
+            Assert.Equal(new long[] { 6 }, t.shape);
+            Assert.Equal(torch.float32, t.dtype);
+            Assert.True(expected.allclose(t));
+        }
+
         [Fact]
         public void TestLoadJIT_Func()
         {
@@ -34,7 +55,7 @@ public void TestLoadJIT_Func()
         }
 
         [Fact]
-        public void TestLoadJIT_1()
+        public void TestLoadJIT_5()
         {
             var input = torch.ones(10);
             var expected = torch.tensor(new float[] { 0.313458264f, 0, 0.9996568f, 0, 0, 0 });
@@ -54,6 +75,39 @@ public void TestLoadJIT_1()
             Assert.True(expected.allclose(t));
         }
 
+        [Fact]
+        public void TestLoadJIT_6()
+        {
+            var input = torch.ones(10);
+            var expected = torch.tensor(new float[] { 0.313458264f, 0, 0.9996568f, 0, 0, 0 });
+
+            // One linear layer followed by ReLU.
+            var m = torch.jit.load<Tensor, Tensor>(@"linrelu.script.dat");
+            if (torch.cuda.is_available()) {
+                m = m.to(torch.CUDA);
+                input = input.to(torch.CUDA);
+                expected = expected.to(torch.CUDA);
+            }
+
+            int i = 0;
+            m.register_forward_pre_hook((m,t) => { i += 1; return t; });
+            m.register_forward_pre_hook((m, t) => { i += 2; return t; });
+            m.register_forward_hook((m, t1, t2) => { i += 4; return t2; });
+            m.register_forward_hook((m, t1, t2) => { i += 8; return t2; });
+
+            var t = m.forward(input);
+
+            Assert.Equal(0, i);
+
+            t = m.call(input);
+
+            Assert.Equal(15, i);
+
+            Assert.Equal(new long[] { 6 }, t.shape);
+            Assert.Equal(torch.float32, t.dtype);
+            Assert.True(expected.allclose(t));
+        }
+
         [Fact]
         public void TestSaveJIT()
         {