Massive updates for TensorFlowOpLayer #652

Author: Oceania2018

src/TensorFlowNET.Core/Functions/ConcreteFunction.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using Tensorflow.Framework.Models;
 using Tensorflow.Graphs;
@@ -11,11 +12,34 @@ namespace Tensorflow.Functions
     /// </summary>
     public class ConcreteFunction : IDisposable
     {
-        public string Name => _handle == IntPtr.Zero ? string.Empty : c_api.StringPiece(c_api.TF_FunctionName(_handle));
         IntPtr _handle;
+        FuncGraph func_graph;
+
+        public string Name
+        {
+            get
+            {
+                if (func_graph != null)
+                    return func_graph.FuncName;
+
+                return _handle == IntPtr.Zero ? string.Empty : c_api.StringPiece(c_api.TF_FunctionName(_handle));
+            }
+        }
+        
         public Tensor[] Outputs;
+        public Type ReturnType;
         public TensorSpec[] OutputStructure;
 
+        public ConcreteFunction(string name)
+        {
+            func_graph = new FuncGraph(name);
+        }
+
+        public ConcreteFunction(FuncGraph graph, Dictionary<string, string> attrs)
+        {
+            func_graph = graph;
+        }
+
         public ConcreteFunction(Func<Tensor, Tensor> func, TF_DataType dtype)
         {
             string func_name = $"autograph_{Guid.NewGuid()}_{func.Method.Name}";
@@ -28,8 +52,8 @@ public ConcreteFunction(Func<Tensor, Tensor> func, TF_DataType dtype)
 
                 var opers = graph._nodes_by_name.Values.Select(x => x as Operation).ToArray();
                 _handle = graph.ToGraph(opers,
-                    new Operation[] { input },
-                    new Operation[] { output },
+                    new[] { input },
+                    new[] { output },
                     null);
             }
         }
@@ -48,8 +72,8 @@ public ConcreteFunction(Func<Tensor, IDatasetV2> func, TF_DataType dtype)
 
                 var opers = graph._nodes_by_name.Values.Select(x => x as Operation).ToArray();
                 _handle = graph.ToGraph(opers,
-                    new Operation[] { input },
-                    new Operation[] { output.variant_tensor.op },
+                    new[] { input },
+                    new[] { output.variant_tensor },
                     null);
             }
         }
@@ -72,12 +96,38 @@ public ConcreteFunction(Func<Tensor, (Tensor, Tensor), (Tensor, Tensor)> func,
 
                 var opers = graph._nodes_by_name.Values.Select(x => x as Operation).ToArray();
                 _handle = graph.ToGraph(opers,
-                    new Operation[] { input1, input2, input3 },
-                    new Operation[] { outputs.Item1.op, outputs.Item2.op },
+                    new[] { input1, input2, input3 },
+                    new[] { outputs.Item1, outputs.Item2 },
                     null);
             }
         }
 
+        public void ToGraph(Tensors inputs, Tensors outputs)
+        {
+            var opers = func_graph._nodes_by_name.Values.Select(x => x as Operation).ToArray();
+            _handle = func_graph.ToGraph(opers,
+                inputs,
+                outputs,
+                null);
+        }
+
+        public Tensors Invoke(Tensors inputs)
+        {
+            var forward_backward = SelectForwardAndBackwardFunctions(inputs, 1, tf.Context.executing_eagerly());
+            var (forward_function, args_with_tangents) = forward_backward.Forward();
+            Tensors flat_outputs = null;
+            if (tf.Context.executing_eagerly())
+                flat_outputs = forward_function.Call(args_with_tangents);
+            forward_backward.Record(flat_outputs);
+            return flat_outputs;
+        }
+
+        ForwardBackwardCall SelectForwardAndBackwardFunctions(Tensors args, int possible_gradient_type, bool executing_eagerly)
+        {
+            var functions = new FirstOrderTapeGradientFunctions(func_graph, false);
+            return new ForwardBackwardCall(functions, args, tape_watching: true);
+        }
+
         public void Dispose()
         {
             c_api.TFE_ContextRemoveFunction(tf.Context.Handle, Name, tf.Status.Handle);

src/TensorFlowNET.Core/Functions/EagerDefinedFunction.cs

@@ -0,0 +1,44 @@
+using Google.Protobuf;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using Tensorflow.Graphs;
+using static Tensorflow.Binding;
+
+namespace Tensorflow.Functions
+{
+    public class EagerDefinedFunction
+    {
+        public int _num_outputs;
+        public string Name => _func_graph.FuncName;
+
+        FuncGraph _func_graph;
+        public EagerDefinedFunction(string name, FuncGraph graph, 
+            Tensors inputs, Tensors outputs,
+            Dictionary<string, string> attrs)
+        {
+            _num_outputs = outputs.Length;
+            
+            var input_ops = inputs.Select(x => x.op).ToArray();
+            var operations = graph.get_operations().Where(x => !input_ops.Contains(x.op))
+                .Select(x => x as Operation).ToArray();
+            var output_names = new string[0];
+
+            _func_graph = new FuncGraph(graph, name, attrs);
+            _func_graph.ToGraph(operations, inputs, outputs, output_names);
+        }
+
+        public Tensors Call(Tensors args)
+        {
+            var results = tf.Runner.TFE_Execute(tf.Context,
+                tf.Context.DeviceName,
+                _func_graph.FuncName,
+                args,
+                null,
+                _num_outputs);
+
+            return results;
+        }
+    }
+}

src/TensorFlowNET.Core/Functions/FirstOrderTapeGradientFunctions.cs

@@ -0,0 +1,25 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Tensorflow.Graphs;
+
+namespace Tensorflow.Functions
+{
+    public class FirstOrderTapeGradientFunctions : TapeGradientFunctions
+    {
+        public FirstOrderTapeGradientFunctions(FuncGraph func_graph,
+            bool need_gradients_for_jvps) : base(func_graph, 
+                need_gradients_for_jvps)
+        {
+
+        }
+
+        public override EagerDefinedFunction ForwardAndBackwardFunctions(Tensors inference_args)
+        {
+            var outputs = _func_graph.Outputs;
+            (_forward, _forward_graph, _backward, _forwardprop_output_indices, _num_forwardprop_outputs) 
+                = BuildFunctionsForOutputs(outputs, inference_args);
+            return _forward;
+        }
+    }
+}

src/TensorFlowNET.Core/Functions/ForwardBackwardCall.cs

@@ -0,0 +1,38 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace Tensorflow.Functions
+{
+    /// <summary>
+    /// Holds the state of a function call between execution and recording.
+    /// </summary>
+    public class ForwardBackwardCall
+    {
+        TapeGradientFunctions _functions;
+        Tensors _inference_args;
+        Tensors _input_tangents;
+        bool _tape_watching;
+
+        public ForwardBackwardCall(TapeGradientFunctions functions, 
+            Tensors inference_args, 
+            bool tape_watching)
+        {
+            _functions = functions;
+            _inference_args = inference_args;
+            _tape_watching = tape_watching;
+        }
+
+        public (EagerDefinedFunction, Tensors) Forward()
+        {
+            var forward_function = _functions.Forward(_inference_args);
+            return (forward_function, _inference_args);
+        }
+
+        public void Record(Tensors flat_outputs)
+        {
+            if (_tape_watching && flat_outputs != null)
+                _functions.Record(flat_outputs, _inference_args);
+        }
+    }
+}

src/TensorFlowNET.Core/Functions/TapeGradientFunctions.cs

@@ -0,0 +1,120 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+using Tensorflow.Graphs;
+using static Tensorflow.Binding;
+using static Tensorflow.tensorflow;
+
+namespace Tensorflow.Functions
+{
+    /// <summary>
+    /// Caches forward and backward functions compatible with eager gradients.
+    /// </summary>
+    public abstract class TapeGradientFunctions
+    {
+        string FORWARD_FUNCTION_ATTRIBUTE_NAME = "forward_function_name";
+        string BACKWARD_FUNCTION_ATTRIBUTE_NAME = "backward_function_name";
+        string _FORWARD_PREFIX = "__forward_";
+        string _BACKWARD_PREFIX = "__backward_";
+        string _INFERENCE_PREFIX = "__inference_";
+
+        protected FuncGraph _func_graph;
+        protected EagerDefinedFunction _forward;
+        protected FuncGraph _forward_graph;
+        protected List<int> _forwardprop_output_indices;
+        protected int _num_forwardprop_outputs;
+        protected ConcreteFunction _backward;
+
+        public TapeGradientFunctions(FuncGraph func_graph,
+            bool need_gradients_for_jvps)
+        {
+            _func_graph = func_graph;
+        }
+
+        public EagerDefinedFunction Forward(Tensors inference_args)
+        {
+            return ForwardAndBackwardFunctions(inference_args);
+        }
+
+        /// <summary>
+        /// Record the function call operation.
+        /// </summary>
+        /// <param name="flat_outputs"></param>
+        /// <param name="inference_args"></param>
+        public void Record(Tensors flat_outputs, Tensors inference_args)
+        {
+            var (backward_function, to_record) = _wrap_backward_function(_forward_graph, _backward, flat_outputs);
+            tf.Runner.RecordGradient(_forward.Name, flat_outputs, new object[0], inference_args, 
+                getBackwardFunction: () => backward_function);
+        }
+
+        (BackwardFunction, Tensors) _wrap_backward_function(FuncGraph forward_graph, ConcreteFunction backward, Tensors flat_outputs)
+        {
+            BackwardFunction _backward_function_wrapper = (output_grads, unneeded_gradients) =>
+            {
+                return new Tensor[0];
+
+                /*var gradients = ops.gradientFunctions[op_name](new EagerOperation
+                {
+                    Name = op_name,
+                    NumInputs = op_inputs.Length,
+                    Inputs = op_inputs,
+                    NumOutputs = op_outputs.Length,
+                    Outputs = op_outputs,
+                    SkipInputIndices = unneeded_gradients,
+                    Attrs = attrs
+                }, output_grads);
+
+                return gradients;*/
+            };
+
+            return (_backward_function_wrapper, flat_outputs);
+        }
+
+        protected (EagerDefinedFunction, FuncGraph, ConcreteFunction, List<int>, int) 
+            BuildFunctionsForOutputs(Tensors outputs, Tensors inference_args)
+        {
+            var trainable_outputs = new List<Tensor>();
+            var trainable_indices = new List<int>();
+            foreach(var (index, output) in enumerate(outputs))
+            {
+                if (gradients_util.IsTrainable(output))
+                {
+                    trainable_outputs.Add(output);
+                    trainable_indices.Add(index);
+                }
+            }
+
+            var gradients_wrt_outputs = new List<Tensor>();
+            var backwards_graph = new FuncGraph($"{_BACKWARD_PREFIX}{_func_graph.FuncName}_{ops.uid()}");
+            foreach (var output in trainable_outputs)
+                gradients_wrt_outputs.Add(tf.placeholder(output.dtype, output.shape));
+            var gradients_wrt_inputs = gradients_util._GradientsHelper(trainable_outputs.ToArray(),
+                _func_graph.Inputs,
+                grad_ys: gradients_wrt_outputs.ToArray(),
+                src_graph: _func_graph);
+
+            tf.Context.restore_mode();
+
+            var forward_function_name = $"{_FORWARD_PREFIX}{_func_graph.FuncName}_{ops.uid()}";
+            var backward_function_attr = new Dictionary<string, string>();
+            backward_function_attr[FORWARD_FUNCTION_ATTRIBUTE_NAME] = forward_function_name;
+            backwards_graph.Inputs = gradients_wrt_outputs;
+            backwards_graph.Outputs = gradients_wrt_inputs;
+
+            var backward_function = new ConcreteFunction(backwards_graph, backward_function_attr);
+
+            var forward_function_attr = new Dictionary<string, string>();
+            forward_function_attr[BACKWARD_FUNCTION_ATTRIBUTE_NAME] = backward_function.Name;
+            var forward_function = new EagerDefinedFunction(forward_function_name, _func_graph, 
+                _func_graph.Inputs, _func_graph.Outputs, forward_function_attr);
+            
+            return (forward_function, _func_graph, backward_function, null, 0);
+        }
+
+        public virtual EagerDefinedFunction ForwardAndBackwardFunctions(Tensors inference_args)
+        {
+            throw new NotImplementedException("");
+        }
+    }
+}

src/TensorFlowNET.Core/Functions/c_api.function.cs

@@ -47,6 +47,9 @@ public static extern IntPtr TF_GraphToFunction(IntPtr fn_body, string fn_name,
             string description,
             SafeStatusHandle status);
 
+        [DllImport(TensorFlowLibName)]
+        public static extern IntPtr TF_FunctionSetAttrValueProto(IntPtr func, string attr_name, byte[] proto, int proto_len, SafeStatusHandle status);
+
         [DllImport(TensorFlowLibName)]
         public static extern IntPtr TF_FunctionName(IntPtr func);
 

src/TensorFlowNET.Core/Gradients/ITape.cs

@@ -13,8 +13,6 @@ public interface ITape
         void RecordOperation(string op_type,
             Tensor[] input_tensors,
             TapeTensor[] output_tensors,
-            long[] input_tensor_id,
-            TF_DataType[] input_dtypes,
             Func<BackwardFunction> backward_function_getter);
 
         void VariableAccessed(ResourceVariable variable);

src/TensorFlowNET.Core/Gradients/Tape.RecordOperation.cs

@@ -3,6 +3,7 @@
 using Tensorflow.Util;
 using static Tensorflow.tensorflow;
 using static Tensorflow.Binding;
+using System.Linq;
 
 namespace Tensorflow.Gradients
 {
@@ -14,18 +15,19 @@ public partial class Tape
         public void RecordOperation(string op_type,
             Tensor[] input_tensors,
             TapeTensor[] output_tensors,
-            long[] input_tensor_id,
-            TF_DataType[] input_dtypes,
             Func<BackwardFunction> backward_function_getter)
         {
-            if (!ShouldRecord(input_tensor_id, input_dtypes))
+            var input_ids = input_tensors.Select(x => x.Id).ToArray();
+            var input_dtypes = input_tensors.Select(x => x.dtype).ToArray();
+
+            if (!ShouldRecord(input_ids, input_dtypes))
             {
                 return;
             }
 
             long op_id = next_op_id_++;
-            var ids = new List<long>(input_tensor_id.Length);
-            foreach (var i in input_tensor_id)
+            var ids = new List<long>(input_ids.Length);
+            foreach (var i in input_ids)
             {
                 tensor_usage_[i]++;
                 ids.Add(i);