Remove ISchema in TensorflowUtils.cs and polish its shape translation

Author: wschin

src/Microsoft.ML.TensorFlow/TensorFlow/TensorflowUtils.cs

@@ -20,51 +20,49 @@ public static class TensorFlowUtils
 
         internal static Schema GetModelSchema(IExceptionContext ectx, TFGraph graph, string opType = null)
         {
-            var res = new List<KeyValuePair<string, ColumnType>>();
-            var opTypeGetters = new List<MetadataUtils.MetadataGetter<ReadOnlyMemory<char>>>();
-            var inputOpsGetters = new List<MetadataUtils.MetadataGetter<VBuffer<ReadOnlyMemory<char>>>>();
-            var inputOpsLengths = new List<int>();
+            var schemaBuilder = new SchemaBuilder();
             foreach (var op in graph)
             {
                 if (opType != null && opType != op.OpType)
                     continue;
+
                 var tfType = op[0].OutputType;
+                // Determine element type in Tensorflow tensor. For example, a vector of floats may get NumberType.R4 here.
                 var mlType = Tf2MlNetTypeOrNull(tfType);
 
                 // If the type is not supported in ML.NET then we cannot represent it as a column in an ISchema.
                 // We also cannot output it with a TensorFlowTransform, so we skip it.
                 if (mlType == null)
                     continue;
 
-                var shape = graph.GetTensorShape(op[0]);
-                var shapeArray = shape.ToIntArray();
-
-                inputOpsLengths.Add(op.NumInputs);
-                MetadataUtils.MetadataGetter<VBuffer<ReadOnlyMemory<char>>> inputOpsGetter = null;
+                // Construct the final ML.NET type of a Tensorflow variable.
+                var shapeArray = graph.GetTensorShape(op[0]).ToIntArray();
+                var columnType = new VectorType(mlType);
+                if (Utils.Size(shapeArray) > 0 && shapeArray.Skip(1).All(x => x > 0))
+                    columnType = new VectorType(mlType, shapeArray[0] > 0 ? shapeArray : shapeArray.Skip(1).ToArray());
+
+                // There can be at most two metadata fields.
+                //  1. The first field always presents. Its value is this operator's ID. For example,
+                //     if an output is produced by an operator named "A", the value of this field should be "A".
+                //  2. The second field stores operators whose outputs are consumed by this operator. In other words,
+                //     these values are names of some upstream operators which should be evaluated before executing
+                //     the current operator. It's possible that one operator doesn't need any input, so this field
+                //     can be missing.
+                var metadataBuilder = new MetadataBuilder();
+                metadataBuilder.Add(OpType, TextType.Instance, (ref ReadOnlyMemory<char> value) => value = op.OpType.AsMemory());
                 if (op.NumInputs > 0)
-                {
-                    var inputOps = new ReadOnlyMemory<char>[op.NumInputs];
-                    for (int i = 0; i < op.NumInputs; i++)
-                    {
-                        var input = op.GetInput(i);
-                        inputOps[i] = new ReadOnlyMemory<char>(input.Operation.Name.ToArray());
-                    }
-                    inputOpsGetter = (int col, ref VBuffer<ReadOnlyMemory<char>> dst) =>
-                        dst = new VBuffer<ReadOnlyMemory<char>>(op.NumInputs, inputOps);
-                }
-                inputOpsGetters.Add(inputOpsGetter);
-
-                MetadataUtils.MetadataGetter<ReadOnlyMemory<char>> opTypeGetter =
-                    (int col, ref ReadOnlyMemory<char> dst) => dst = new ReadOnlyMemory<char>(op.OpType.ToArray());
-                opTypeGetters.Add(opTypeGetter);
-
-                var columnType = Utils.Size(shapeArray) == 1 && shapeArray[0] <= 0 ? new VectorType(mlType) :
-                    Utils.Size(shapeArray) > 0 && shapeArray.Skip(1).All(x => x > 0) ?
-                        new VectorType(mlType, shapeArray[0] > 0 ? shapeArray : shapeArray.Skip(1).ToArray())
-                        : new VectorType(mlType);
-                res.Add(new KeyValuePair<string, ColumnType>(op.Name, columnType));
+                    metadataBuilder.Add(InputOps, new VectorType(TextType.Instance, op.NumInputs),
+                        (ref VBuffer<ReadOnlyMemory<char>> value) =>
+                            {
+                                var bufferEditor = VBufferEditor.Create(ref value, op.NumInputs);
+                                for (int i = 0; i < op.NumInputs; ++i)
+                                    bufferEditor.Values[i] = op.GetInput(i).Operation.Name.AsMemory();
+                                value = bufferEditor.Commit();
+                            });
+
+                schemaBuilder.AddColumn(op.Name, columnType, metadataBuilder.GetMetadata());
             }
-            return Schema.Create(new TensorFlowSchema(ectx, res.ToArray(), opTypeGetters.ToArray(), inputOpsGetters.ToArray(), inputOpsLengths.ToArray()));
+            return schemaBuilder.GetSchema();
         }
 
         /// <summary>
@@ -75,25 +73,27 @@ internal static Schema GetModelSchema(IExceptionContext ectx, TFGraph graph, str
         /// of kind <see cref="OpType"/>, indicating the operation type of the node, and if that node has inputs in the graph,
         /// it contains metadata of kind <see cref="InputOps"/>, indicating the names of the input nodes.
         /// </summary>
-        /// <param name="env">The environment to use.</param>
-        /// <param name="modelPath">Model to load.</param>
-        public static Schema GetModelSchema(IHostEnvironment env, string modelPath)
+        /// <param name="ectx">An <see cref="IExceptionContext"/>.</param>
+        /// <param name="modelFile">The name of the file containing the TensorFlow model. Currently only frozen model
+        /// format is supported.</param>
+        public static Schema GetModelSchema(IExceptionContext ectx, string modelFile)
         {
-            var model = LoadTensorFlowModel(env, modelPath);
-            return GetModelSchema(env, model.Session.Graph);
+            var bytes = File.ReadAllBytes(modelFile);
+            var session = LoadTFSession(ectx, bytes, modelFile);
+            return GetModelSchema(ectx, session.Graph);
         }
 
         /// <summary>
         /// This is a convenience method for iterating over the nodes of a TensorFlow model graph. It
-        /// iterates over the columns of the <see cref="ISchema"/> returned by <see cref="GetModelSchema(IHostEnvironment, string)"/>,
+        /// iterates over the columns of the <see cref="ISchema"/> returned by <see cref="GetModelSchema(IExceptionContext, string)"/>,
         /// and for each one it returns a tuple containing the name, operation type, column type and an array of input node names.
         /// This method is convenient for filtering nodes based on certain criteria, for example, by the operation type.
         /// </summary>
-        /// <param name="modelPath">Model to load.</param>
+        /// <param name="modelFile"></param>
         /// <returns></returns>
-        public static IEnumerable<(string, string, ColumnType, string[])> GetModelNodes(string modelPath)
+        public static IEnumerable<(string, string, ColumnType, string[])> GetModelNodes(string modelFile)
         {
-            var schema = GetModelSchema(new MLContext(), modelPath);
+            var schema = GetModelSchema(null, modelFile);
 
             for (int i = 0; i < schema.Count; i++)
             {
@@ -308,12 +308,6 @@ private static void CreateTempDirectoryWithAcl(string folder, string identity)
             }
         }
 
-        /// <summary>
-        /// Load TensorFlow model into memory.
-        /// </summary>
-        /// <param name="env">The environment to use.</param>
-        /// <param name="modelPath">The model to load.</param>
-        /// <returns></returns>
         public static TensorFlowModelInfo LoadTensorFlowModel(IHostEnvironment env, string modelPath)
         {
             var session = GetSession(env, modelPath);
@@ -358,55 +352,5 @@ internal static bool IsTypeSupported(TFDataType tfoutput)
                     return false;
             }
         }
-
-        private sealed class TensorFlowSchema : SimpleSchemaBase
-        {
-            private readonly MetadataUtils.MetadataGetter<ReadOnlyMemory<char>>[] _opTypeGetters;
-            private readonly MetadataUtils.MetadataGetter<VBuffer<ReadOnlyMemory<char>>>[] _inputOpsGetters;
-            private readonly int[] _inputOpsLengths;
-
-            public TensorFlowSchema(IExceptionContext ectx, KeyValuePair<string, ColumnType>[] columns,
-                MetadataUtils.MetadataGetter<ReadOnlyMemory<char>>[] opTypeGetters,
-                MetadataUtils.MetadataGetter<VBuffer<ReadOnlyMemory<char>>>[] inputOpsGetters, int[] inputOpsLengths)
-                : base(ectx, columns)
-            {
-                ectx.CheckParam(Utils.Size(opTypeGetters) == ColumnCount, nameof(opTypeGetters));
-                ectx.CheckParam(Utils.Size(inputOpsGetters) == ColumnCount, nameof(inputOpsGetters));
-                ectx.CheckParam(Utils.Size(inputOpsLengths) == ColumnCount, nameof(inputOpsLengths));
-
-                _opTypeGetters = opTypeGetters;
-                _inputOpsGetters = inputOpsGetters;
-                _inputOpsLengths = inputOpsLengths;
-            }
-
-            protected override void GetMetadataCore<TValue>(string kind, int col, ref TValue value)
-            {
-                Ectx.Assert(0 <= col && col < ColumnCount);
-                if (kind == OpType)
-                    _opTypeGetters[col].Marshal(col, ref value);
-                else if (kind == InputOps && _inputOpsGetters[col] != null)
-                    _inputOpsGetters[col].Marshal(col, ref value);
-                else
-                    throw Ectx.ExceptGetMetadata();
-            }
-
-            protected override ColumnType GetMetadataTypeOrNullCore(string kind, int col)
-            {
-                Ectx.Assert(0 <= col && col < ColumnCount);
-                if (kind == OpType)
-                    return TextType.Instance;
-                if (kind == InputOps && _inputOpsGetters[col] != null)
-                    return new VectorType(TextType.Instance, _inputOpsLengths[col]);
-                return null;
-            }
-
-            protected override IEnumerable<KeyValuePair<string, ColumnType>> GetMetadataTypesCore(int col)
-            {
-                Ectx.Assert(0 <= col && col < ColumnCount);
-                yield return new KeyValuePair<string, ColumnType>(OpType, TextType.Instance);
-                if (_inputOpsGetters[col] != null)
-                    yield return new KeyValuePair<string, ColumnType>(InputOps, new VectorType(TextType.Instance, _inputOpsLengths[col]));
-            }
-        }
     }
 }