Exposed TensorFLow session as TensorFlowModelInfo class (#1191)

Author: zeahmed

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

@@ -4,7 +4,6 @@
 
 using Microsoft.ML.Runtime;
 using Microsoft.ML.Runtime.Data;
-using Microsoft.ML.Runtime.ImageAnalytics.EntryPoints;
 using Microsoft.ML.Runtime.Internal.Utilities;
 using System;
 using System.Collections.Generic;
@@ -21,14 +20,16 @@ public static class TensorFlowUtils
         public const string OpType = "OpType";
         public const string InputOps = "InputOps";
 
-        private static ISchema GetModelSchema(IExceptionContext ectx, TFGraph graph)
+        internal static ISchema 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>();
             foreach (var op in graph)
             {
+                if (opType != null && opType != op.OpType)
+                    continue;
                 var tfType = op[0].OutputType;
                 var mlType = Tf2MlNetTypeOrNull(tfType);
 
@@ -55,12 +56,11 @@ private static ISchema GetModelSchema(IExceptionContext ectx, TFGraph graph)
                 }
                 inputOpsGetters.Add(inputOpsGetter);
 
-                var opType = op.OpType;
                 MetadataUtils.MetadataGetter<ReadOnlyMemory<char>> opTypeGetter =
-                    (int col, ref ReadOnlyMemory<char> dst) => dst = new ReadOnlyMemory<char>(opType.ToArray());
+                    (int col, ref ReadOnlyMemory<char> dst) => dst = new ReadOnlyMemory<char>(op.OpType.ToArray());
                 opTypeGetters.Add(opTypeGetter);
 
-                var columnType = Utils.Size(shapeArray) == 1 && shapeArray[0] == -1 ? new VectorType(mlType) :
+                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);
@@ -308,6 +308,12 @@ private static void CreateTempDirectoryWithAcl(string folder, string identity)
             }
         }
 
+        public static TensorFlowModelInfo LoadTensorFlowModel(IHostEnvironment env, string modelPath)
+        {
+            var session = GetSession(env, modelPath);
+            return new TensorFlowModelInfo(env, session, modelPath);
+        }
+
         internal static TFSession GetSession(IHostEnvironment env, string modelPath)
         {
             Contracts.Check(env != null, nameof(env));

src/Microsoft.ML.TensorFlow/TensorFlowModelInfo.cs

@@ -0,0 +1,60 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.ML.Runtime;
+using Microsoft.ML.Runtime.Data;
+using Microsoft.ML.Transforms.TensorFlow;
+
+namespace Microsoft.ML.Transforms
+{
+    /// <summary>
+    /// This class holds the information related to TensorFlow model and session.
+    /// It provides a convenient way to query model schema as follows.
+    /// <list type="bullet">
+    ///    <item>
+    ///      <description>Get complete schema by calling <see cref="GetModelSchema()"/>.</description>
+    ///    </item>
+    ///    <item>
+    ///      <description>Get schema related to model input(s) by calling <see cref="GetInputSchema()"/>.</description>
+    ///    </item>
+    /// </list>
+    /// </summary>
+    public class TensorFlowModelInfo
+    {
+        internal TFSession Session { get; }
+        public string ModelPath { get; }
+
+        private readonly IHostEnvironment _env;
+
+        /// <summary>
+        /// Instantiates <see cref="TensorFlowModelInfo"/>.
+        /// </summary>
+        /// <param name="env">An <see cref="IHostEnvironment"/> object.</param>
+        /// <param name="session">TensorFlow session object.</param>
+        /// <param name="modelLocation">Location of the model from where <paramref name="session"/> was loaded.</param>
+        internal TensorFlowModelInfo(IHostEnvironment env, TFSession session, string modelLocation)
+        {
+            Session = session;
+            ModelPath = modelLocation;
+            _env = env;
+        }
+
+        /// <summary>
+        /// Get <see cref="ISchema"/> for complete model. Every node in the TensorFlow model will be included in the <see cref="ISchema"/> object.
+        /// </summary>
+        public ISchema GetModelSchema()
+        {
+            return TensorFlowUtils.GetModelSchema(_env, Session.Graph);
+        }
+
+        /// <summary>
+        /// Get <see cref="ISchema"/> for only those nodes which are marked "Placeholder" in the TensorFlow model.
+        /// This method is convenient for exploring the model input(s) in case TensorFlow graph is very large.
+        /// </summary>
+        public ISchema GetInputSchema()
+        {
+            return TensorFlowUtils.GetModelSchema(_env, Session.Graph, "Placeholder");
+        }
+    }
+}

src/Microsoft.ML.TensorFlow/TensorflowTransform.cs

@@ -189,7 +189,9 @@ private static VersionInfo GetVersionInfo()
         }
 
         /// <summary>
-        /// Convenience constructor for public facing API.
+        /// Creates <see cref="IDataTransform"/> using <see cref="TensorFlowTransform"/>.
+        /// This convenience method get the model file as input and loads the model internally.
+        /// If the model is already loaded please <see cref="TensorFlowTransform.Create(IHostEnvironment, IDataView, TensorFlowModelInfo, string[], string[])"/> to avoid reloading of model.
         /// </summary>
         /// <param name="env">Host Environment.</param>
         /// <param name="input">Input <see cref="IDataView"/>. This is the output from previous transform or loader.</param>
@@ -201,6 +203,21 @@ public static IDataTransform Create(IHostEnvironment env, IDataView input, strin
             return new TensorFlowTransform(env, TensorFlowUtils.GetSession(env, model), source, names, TensorFlowUtils.IsSavedModel(env, model) ? model : null, false).MakeDataTransform(input);
         }
 
+        /// <summary>
+        /// Creates <see cref="IDataTransform"/> using <see cref="TensorFlowTransform"/>.
+        /// This convenience method avoids reloading of TensorFlow model.
+        /// It is useful in a situation where user has already loaded TensorFlow model using <see cref="TensorFlowUtils.LoadTensorFlowModel(IHostEnvironment, string)"/> for inspecting model schema.
+        /// </summary>
+        /// <param name="env">Host Environment.</param>
+        /// <param name="input">Input <see cref="IDataView"/>. This is the output from previous transform or loader.</param>
+        /// <param name="tfModelInfo"> <see cref="TensorFlowModelInfo"/> object created with <see cref="TensorFlowUtils.LoadTensorFlowModel(IHostEnvironment, string)"/>.</param>
+        /// <param name="names">Name of the output column(s). Keep it same as in the Tensorflow model.</param>
+        /// <param name="source">Name of the input column(s). Keep it same as in the Tensorflow model.</param>
+        public static IDataTransform Create(IHostEnvironment env, IDataView input, TensorFlowModelInfo tfModelInfo, string[] names, string[] source)
+        {
+            return new TensorFlowTransform(env, tfModelInfo.Session, source, names, TensorFlowUtils.IsSavedModel(env, tfModelInfo.ModelPath) ? tfModelInfo.ModelPath : null, false).MakeDataTransform(input);
+        }
+
         // Factory method for SignatureLoadModel.
         private static TensorFlowTransform Create(IHostEnvironment env, ModelLoadContext ctx)
         {
@@ -1085,6 +1102,11 @@ public TensorFlowEstimator(IHostEnvironment env, string model, string[] inputs,
         {
         }
 
+        public TensorFlowEstimator(IHostEnvironment env, TensorFlowModelInfo tensorFlowModel, string[] inputs, string[] outputs)
+           : this(env, new TensorFlowTransform(env, tensorFlowModel.Session, inputs, outputs, TensorFlowUtils.IsSavedModel(env, tensorFlowModel.ModelPath) ? tensorFlowModel.ModelPath : null, false))
+        {
+        }
+
         public TensorFlowEstimator(IHostEnvironment env, TensorFlowTransform transformer)
             : base(Contracts.CheckRef(env, nameof(env)).Register(nameof(TensorFlowTransform)), transformer)
         {
@@ -1127,16 +1149,32 @@ public OutColumn(Vector<float> input, string modelFile)
             {
                 Input = input;
             }
+
+            public OutColumn(Vector<float> input, TensorFlowModelInfo tensorFlowModel)
+                : base(new Reconciler(tensorFlowModel), input)
+            {
+                Input = input;
+            }
         }
 
         private sealed class Reconciler : EstimatorReconciler
         {
             private readonly string _modelFile;
+            private readonly TensorFlowModelInfo _tensorFlowModel;
 
             public Reconciler(string modelFile)
             {
                 Contracts.AssertNonEmpty(modelFile);
                 _modelFile = modelFile;
+                _tensorFlowModel = null;
+            }
+
+            public Reconciler(TensorFlowModelInfo tensorFlowModel)
+            {
+                Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));
+
+                _modelFile = null;
+                _tensorFlowModel = tensorFlowModel;
             }
 
             public override IEstimator<ITransformer> Reconcile(IHostEnvironment env,
@@ -1148,15 +1186,22 @@ public override IEstimator<ITransformer> Reconcile(IHostEnvironment env,
                 Contracts.Assert(toOutput.Length == 1);
 
                 var outCol = (OutColumn)toOutput[0];
-                return new TensorFlowEstimator(env, _modelFile, new[] { inputNames[outCol.Input] }, new[] { outputNames[outCol] });
+                if (_modelFile == null)
+                {
+                    return new TensorFlowEstimator(env, _tensorFlowModel, new[] { inputNames[outCol.Input] }, new[] { outputNames[outCol] });
+                }
+                else
+                {
+                    return new TensorFlowEstimator(env, _modelFile, new[] { inputNames[outCol.Input] }, new[] { outputNames[outCol] });
+                }
             }
         }
 
         // REVIEW: this method only covers one use case of using TensorFlow models: consuming one
         // input and producing one output of floats.
         // We could consider selectively adding some more extensions to enable common scenarios.
         /// <summary>
-        /// Run a TensorFlow model on the input column and extract one output column.
+        /// Load the TensorFlow model from <paramref name="modelFile"/> and run it on the input column and extract one output column.
         /// The inputs and outputs are matched to TensorFlow graph nodes by name.
         /// </summary>
         public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, string modelFile)
@@ -1165,5 +1210,16 @@ public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, strin
             Contracts.CheckNonEmpty(modelFile, nameof(modelFile));
             return new OutColumn(input, modelFile);
         }
+
+        /// <summary>
+        /// Run a TensorFlow model provided through <paramref name="tensorFlowModel"/> on the input column and extract one output column.
+        /// The inputs and outputs are matched to TensorFlow graph nodes by name.
+        /// </summary>
+        public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, TensorFlowModelInfo tensorFlowModel)
+        {
+            Contracts.CheckValue(input, nameof(input));
+            Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));
+            return new OutColumn(input, tensorFlowModel);
+        }
     }
 }

test/Microsoft.ML.Tests/ScenariosWithDirectInstantiation/TensorflowTests.cs

@@ -735,8 +735,13 @@ public void TensorFlowTransformCifar()
 
             using (var env = new ConsoleEnvironment())
             {
-                var imageHeight = 32;
-                var imageWidth = 32;
+                var tensorFlowModel = TensorFlowUtils.LoadTensorFlowModel(env, model_location);
+                var schema = tensorFlowModel.GetInputSchema();
+                Assert.True(schema.TryGetColumnIndex("Input", out int column));
+                var type = schema.GetColumnType(column).AsVector;
+                var imageHeight = type.GetDim(0);
+                var imageWidth = type.GetDim(1);
+
                 var dataFile = GetDataPath("images/images.tsv");
                 var imageFolder = Path.GetDirectoryName(dataFile);
                 var data = TextLoader.Create(env, new TextLoader.Arguments()
@@ -770,7 +775,7 @@ public void TensorFlowTransformCifar()
                 }, cropped);
 
 
-                IDataView trans = TensorFlowTransform.Create(env, pixels, model_location, new[] { "Output" }, new[] { "Input" });
+                IDataView trans = TensorFlowTransform.Create(env, pixels, tensorFlowModel, new[] { "Output" }, new[] { "Input" });
 
                 trans.Schema.TryGetColumnIndex("Output", out int output);
                 using (var cursor = trans.GetRowCursor(col => col == output))
@@ -796,8 +801,13 @@ public void TensorFlowTransformCifarSavedModel()
 
             using (var env = new ConsoleEnvironment())
             {
-                var imageHeight = 32;
-                var imageWidth = 32;
+                var tensorFlowModel = TensorFlowUtils.LoadTensorFlowModel(env, model_location);
+                var schema = tensorFlowModel.GetInputSchema();
+                Assert.True(schema.TryGetColumnIndex("Input", out int column));
+                var type = schema.GetColumnType(column).AsVector;
+                var imageHeight = type.GetDim(0);
+                var imageWidth = type.GetDim(1);
+
                 var dataFile = GetDataPath("images/images.tsv");
                 var imageFolder = Path.GetDirectoryName(dataFile);
                 var data = TextLoader.Create(env, new TextLoader.Arguments()
@@ -831,7 +841,7 @@ public void TensorFlowTransformCifarSavedModel()
                 }, cropped);
 
 
-                IDataView trans = TensorFlowTransform.Create(env, pixels, model_location, new[] { "Output" }, new[] { "Input" });
+                IDataView trans = TensorFlowTransform.Create(env, pixels, tensorFlowModel, new[] { "Output" }, new[] { "Input" });
 
                 trans.Schema.TryGetColumnIndex("Output", out int output);
                 using (var cursor = trans.GetRowCursor(col => col == output))

test/Microsoft.ML.Tests/TensorFlowEstimatorTests.cs

@@ -10,6 +10,7 @@
 using Microsoft.ML.Runtime.RunTests;
 using Microsoft.ML.Runtime.Tools;
 using Microsoft.ML.Transforms;
+using Microsoft.ML.Transforms.TensorFlow;
 using System;
 using System.Collections.Generic;
 using System.IO;
@@ -182,6 +183,54 @@ public void TestTensorFlowStatic()
             }
         }
 
+        [Fact]
+        public void TestTensorFlowStaticWithSchema()
+        {
+            var modelLocation = "cifar_model/frozen_model.pb";
+
+            using (var env = new ConsoleEnvironment())
+            {
+                var tensorFlowModel = TensorFlowUtils.LoadTensorFlowModel(env, modelLocation);
+                var schema = tensorFlowModel.GetInputSchema();
+                Assert.True(schema.TryGetColumnIndex("Input", out int column));
+                var type = schema.GetColumnType(column).AsVector;
+                var imageHeight = type.GetDim(0);
+                var imageWidth = type.GetDim(1);
+
+                var dataFile = GetDataPath("images/images.tsv");
+                var imageFolder = Path.GetDirectoryName(dataFile);
+
+                var data = TextLoader.CreateReader(env, ctx => (
+                    imagePath: ctx.LoadText(0),
+                    name: ctx.LoadText(1)))
+                    .Read(new MultiFileSource(dataFile));
+
+                // Note that CamelCase column names are there to match the TF graph node names.
+                var pipe = data.MakeNewEstimator()
+                    .Append(row => (
+                        row.name,
+                        Input: row.imagePath.LoadAsImage(imageFolder).Resize(imageHeight, imageWidth).ExtractPixels(interleaveArgb: true)))
+                    .Append(row => (row.name, Output: row.Input.ApplyTensorFlowGraph(tensorFlowModel)));
+
+                TestEstimatorCore(pipe.AsDynamic, data.AsDynamic);
+
+                var result = pipe.Fit(data).Transform(data).AsDynamic;
+                result.Schema.TryGetColumnIndex("Output", out int output);
+                using (var cursor = result.GetRowCursor(col => col == output))
+                {
+                    var buffer = default(VBuffer<float>);
+                    var getter = cursor.GetGetter<VBuffer<float>>(output);
+                    var numRows = 0;
+                    while (cursor.MoveNext())
+                    {
+                        getter(ref buffer);
+                        Assert.Equal(10, buffer.Length);
+                        numRows += 1;
+                    }
+                    Assert.Equal(3, numRows);
+                }
+            }
+        }
 
         private void ValidateTensorFlowTransformer(IDataView result)
         {