Internalization of TensorFlowUtils.cs and refactored TensorFlowCatalog.

docs/samples/Microsoft.ML.Samples/Dynamic/TensorFlow/TextClassification.cs

@@ -45,7 +45,7 @@ public static void Example()
             // Load the TensorFlow model once.
             //      - Use it for quering the schema for input and output in the model
             //      - Use it for prediction in the pipeline.
-            var modelInfo = TensorFlowUtils.LoadTensorFlowModel(mlContext, modelLocation);
+            var modelInfo = mlContext.Model.LoadTensorFlowModel(modelLocation);
             var schema = modelInfo.GetModelSchema();
             var featuresType = (VectorType)schema["Features"].Type;
             Console.WriteLine("Name: {0}, Type: {1}, Shape: (-1, {2})", "Features", featuresType.ItemType.RawType, featuresType.Dimensions[0]);

src/Microsoft.ML.DnnAnalyzer/Microsoft.ML.DnnAnalyzer/DnnAnalyzer.cs

@@ -3,7 +3,10 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
-using Microsoft.ML.Transforms.TensorFlow;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Data.DataView;
+using Microsoft.ML.Data;
 
 namespace Microsoft.ML.DnnAnalyzer
 {
@@ -17,11 +20,40 @@ public static void Main(string[] args)
                 return;
             }
 
-            foreach (var (name, opType, type, inputs) in TensorFlowUtils.GetModelNodes(new MLContext(), args[0]))
+            foreach (var (name, opType, type, inputs) in GetModelNodes(args[0]))
             {
                 var inputsString = inputs.Length == 0 ? "" : $", input nodes: {string.Join(", ", inputs)}";
                 Console.WriteLine($"Graph node: '{name}', operation type: '{opType}', output type: '{type}'{inputsString}");
             }
         }
+
+        private static IEnumerable<(string, string, DataViewType, string[])> GetModelNodes(string modelPath)
+        {
+            var mlContext = new MLContext();
+            var tensorFlowModel = mlContext.Model.LoadTensorFlowModel(modelPath);
+            var schema = tensorFlowModel.GetModelSchema();
+
+            for (int i = 0; i < schema.Count; i++)
+            {
+                var name = schema[i].Name;
+                var type = schema[i].Type;
+
+                var metadataType = schema[i].Metadata.Schema.GetColumnOrNull("TensorflowOperatorType")?.Type;
+                ReadOnlyMemory<char> opType = default;
+                schema[i].Metadata.GetValue("TensorflowOperatorType", ref opType);
+                metadataType = schema[i].Metadata.Schema.GetColumnOrNull("TensorflowUpstreamOperators")?.Type;
+                VBuffer <ReadOnlyMemory<char>> inputOps = default;
+                if (metadataType != null)
+                {
+                    schema[i].Metadata.GetValue("TensorflowUpstreamOperators", ref inputOps);
+                }
+
+                string[] inputOpsResult = inputOps.DenseValues()
+                    .Select(input => input.ToString())
+                    .ToArray();
+
+                yield return (name, opType.ToString(), type, inputOpsResult);
+            }
+        }
     }
 }

src/Microsoft.ML.TensorFlow.StaticPipe/TensorFlowStaticExtensions.cs

@@ -20,7 +20,7 @@ public OutColumn(Vector<float> input, string modelFile)
                 Input = input;
             }
 
-            public OutColumn(Vector<float> input, TensorFlowModelInfo tensorFlowModel)
+            public OutColumn(Vector<float> input, TensorFlowModel tensorFlowModel)
                 : base(new Reconciler(tensorFlowModel), input)
             {
                 Input = input;
@@ -30,7 +30,7 @@ public OutColumn(Vector<float> input, TensorFlowModelInfo tensorFlowModel)
         private sealed class Reconciler : EstimatorReconciler
         {
             private readonly string _modelFile;
-            private readonly TensorFlowModelInfo _tensorFlowModel;
+            private readonly TensorFlowModel _tensorFlowModel;
 
             public Reconciler(string modelFile)
             {
@@ -39,7 +39,7 @@ public Reconciler(string modelFile)
                 _tensorFlowModel = null;
             }
 
-            public Reconciler(TensorFlowModelInfo tensorFlowModel)
+            public Reconciler(TensorFlowModel tensorFlowModel)
             {
                 Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));
 
@@ -81,7 +81,7 @@ public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, strin
         /// 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)
+        public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, TensorFlowModel tensorFlowModel)
         {
             Contracts.CheckValue(input, nameof(input));
             Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));

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

@@ -21,12 +21,12 @@ public static class TensorFlowUtils
         /// Key to access operator's type (a string) in <see cref="DataViewSchema.Column.Metadata"/>.
         /// Its value describes the Tensorflow operator that produces this <see cref="DataViewSchema.Column"/>.
         /// </summary>
-        public const string TensorflowOperatorTypeKind = "TensorflowOperatorType";
+        internal const string TensorflowOperatorTypeKind = "TensorflowOperatorType";
         /// <summary>
         /// Key to access upstream operators' names (a string array) in <see cref="DataViewSchema.Column.Metadata"/>.
         /// Its value states operators that the associated <see cref="DataViewSchema.Column"/>'s generator depends on.
         /// </summary>
-        public const string TensorflowUpstreamOperatorsKind = "TensorflowUpstreamOperators";
+        internal const string TensorflowUpstreamOperatorsKind = "TensorflowUpstreamOperators";
 
         internal static DataViewSchema GetModelSchema(IExceptionContext ectx, TFGraph graph, string opType = null)
         {
@@ -94,50 +94,12 @@ internal static DataViewSchema GetModelSchema(IExceptionContext ectx, TFGraph gr
         /// </summary>
         /// <param name="env">The environment to use.</param>
         /// <param name="modelPath">Model to load.</param>
-        public static DataViewSchema GetModelSchema(IHostEnvironment env, string modelPath)
+        internal static DataViewSchema GetModelSchema(IHostEnvironment env, string modelPath)
         {
             var model = LoadTensorFlowModel(env, modelPath);
             return GetModelSchema(env, model.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="DataViewSchema"/> returned by <see cref="GetModelSchema(IHostEnvironment, 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="env">The environment to use.</param>
-        /// <param name="modelPath">Model to load.</param>
-        /// <returns></returns>
-        public static IEnumerable<(string, string, DataViewType, string[])> GetModelNodes(IHostEnvironment env, string modelPath)
-        {
-            var schema = GetModelSchema(env, modelPath);
-
-            for (int i = 0; i < schema.Count; i++)
-            {
-                var name = schema[i].Name;
-                var type = schema[i].Type;
-
-                var metadataType = schema[i].Metadata.Schema.GetColumnOrNull(TensorflowOperatorTypeKind)?.Type;
-                Contracts.Assert(metadataType != null && metadataType is TextDataViewType);
-                ReadOnlyMemory<char> opType = default;
-                schema[i].Metadata.GetValue(TensorflowOperatorTypeKind, ref opType);
-                metadataType = schema[i].Metadata.Schema.GetColumnOrNull(TensorflowUpstreamOperatorsKind)?.Type;
-                VBuffer<ReadOnlyMemory<char>> inputOps = default;
-                if (metadataType != null)
-                {
-                    Contracts.Assert(metadataType.IsKnownSizeVector() && metadataType.GetItemType() is TextDataViewType);
-                    schema[i].Metadata.GetValue(TensorflowUpstreamOperatorsKind, ref inputOps);
-                }
-
-                string[] inputOpsResult = inputOps.DenseValues()
-                    .Select(input => input.ToString())
-                    .ToArray();
-
-                yield return (name, opType.ToString(), type, inputOpsResult);
-            }
-        }
-
         internal static PrimitiveDataViewType Tf2MlNetType(TFDataType type)
         {
             var mlNetType = Tf2MlNetTypeOrNull(type);
@@ -338,10 +300,10 @@ private static void CreateTempDirectoryWithAcl(string folder, string identity)
         /// <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)
+        internal static TensorFlowModel LoadTensorFlowModel(IHostEnvironment env, string modelPath)
         {
             var session = GetSession(env, modelPath);
-            return new TensorFlowModelInfo(env, session, modelPath);
+            return new TensorFlowModel(env, session, modelPath);
         }
 
         internal static TFSession GetSession(IHostEnvironment env, string modelPath)

src/Microsoft.ML.TensorFlow/TensorFlowModelInfo.cs → src/Microsoft.ML.TensorFlow/TensorFlowModel.cs

@@ -3,7 +3,6 @@
 // See the LICENSE file in the project root for more information.
 
 using Microsoft.Data.DataView;
-using Microsoft.ML.Data;
 using Microsoft.ML.Transforms.TensorFlow;
 
 namespace Microsoft.ML.Transforms
@@ -20,20 +19,20 @@ namespace Microsoft.ML.Transforms
     ///    </item>
     /// </list>
     /// </summary>
-    public class TensorFlowModelInfo
+    public sealed class TensorFlowModel
     {
         internal TFSession Session { get; }
-        public string ModelPath { get; }
+        internal string ModelPath { get; }
 
         private readonly IHostEnvironment _env;
 
         /// <summary>
-        /// Instantiates <see cref="TensorFlowModelInfo"/>.
+        /// Instantiates <see cref="TensorFlowModel"/>.
         /// </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)
+        internal TensorFlowModel(IHostEnvironment env, TFSession session, string modelLocation)
         {
             Session = session;
             ModelPath = modelLocation;

src/Microsoft.ML.TensorFlow/TensorflowCatalog.cs

@@ -5,6 +5,7 @@
 using Microsoft.Data.DataView;
 using Microsoft.ML.Data;
 using Microsoft.ML.Transforms;
+using Microsoft.ML.Transforms.TensorFlow;
 
 namespace Microsoft.ML
 {
@@ -59,7 +60,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// <param name="inputColumnName"> The name of the model input.</param>
         /// <param name="outputColumnName">The name of the requested model output.</param>
         public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog catalog,
-            TensorFlowModelInfo tensorFlowModel,
+            TensorFlowModel tensorFlowModel,
             string outputColumnName,
             string inputColumnName)
             => new TensorFlowEstimator(CatalogUtils.GetEnvironment(catalog), new[] { outputColumnName }, new[] { inputColumnName }, tensorFlowModel);
@@ -79,7 +80,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// </format>
         /// </example>
         public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog catalog,
-            TensorFlowModelInfo tensorFlowModel,
+            TensorFlowModel tensorFlowModel,
             string[] outputColumnNames,
             string[] inputColumnNames)
             => new TensorFlowEstimator(CatalogUtils.GetEnvironment(catalog), outputColumnNames, inputColumnNames, tensorFlowModel);
@@ -102,7 +103,16 @@ public static TensorFlowEstimator TensorFlow(this TransformsCatalog catalog,
         /// <param name="tensorFlowModel">The pre-loaded TensorFlow model.</param>
         public static TensorFlowEstimator TensorFlow(this TransformsCatalog catalog,
             TensorFlowEstimator.Options options,
-            TensorFlowModelInfo tensorFlowModel)
+            TensorFlowModel tensorFlowModel)
             => new TensorFlowEstimator(CatalogUtils.GetEnvironment(catalog), options, tensorFlowModel);
+
+        /// <summary>
+        /// Load TensorFlow model into memory. This is the convenience method that allows the model to be loaded once and subsequently use it for querying schema and creation of
+        /// <see cref="TensorFlowEstimator"/> using <see cref="TensorFlow(TransformsCatalog, TensorFlowEstimator.Options, TensorFlowModel)"/>.
+        /// </summary>
+        /// <param name="catalog">The transform's catalog.</param>
+        /// <param name="modelLocation">Location of the TensorFlow model.</param>
+        public static TensorFlowModel LoadTensorFlowModel(this ModelOperationsCatalog catalog, string modelLocation)
+            => TensorFlowUtils.LoadTensorFlowModel(CatalogUtils.GetEnvironment(catalog), modelLocation);
     }
 }