Internalization of TensorFlowUtils.cs and refactored TensorFlowCatalog.

Author: zeahmed

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

@@ -20,7 +20,7 @@ public static void Example()
             var idv = mlContext.Data.ReadFromEnumerable(data);
 
             // Create a ML pipeline.
-            var pipeline = mlContext.Transforms.ScoreTensorFlowModel(
+            var pipeline = mlContext.Transforms.TensorFlow.ScoreTensorFlowModel(
                 modelLocation, 
                 new[] { nameof(OutputScores.output) },
                 new[] { nameof(TensorData.input) });

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.Transforms.TensorFlow.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]);
@@ -72,7 +72,7 @@ public static void Example()
             var engine = mlContext.Transforms.Text.TokenizeWords("TokenizedWords", "Sentiment_Text")
                 .Append(mlContext.Transforms.Conversion.ValueMap(lookupMap, "Words", "Ids", new[] { ("VariableLenghtFeatures", "TokenizedWords") }))
                 .Append(mlContext.Transforms.CustomMapping(ResizeFeaturesAction, "Resize"))
-                .Append(mlContext.Transforms.ScoreTensorFlowModel(modelInfo, new[] { "Prediction/Softmax" }, new[] { "Features" }))
+                .Append(mlContext.Transforms.TensorFlow.ScoreTensorFlowModel(modelInfo, new[] { "Prediction/Softmax" }, new[] { "Features" }))
                 .Append(mlContext.Transforms.CopyColumns(("Prediction", "Prediction/Softmax")))
                 .Fit(dataView)
                 .CreatePredictionEngine<IMDBSentiment, OutputScores>(mlContext);

src/Microsoft.ML.Data/Transforms/TransformsCatalog.cs

@@ -37,6 +37,11 @@ public sealed class TransformsCatalog
         /// </summary>
         public FeatureSelectionTransforms FeatureSelection { get; }
 
+        /// <summary>
+        /// List of operations for using TensorFlow model.
+        /// </summary>
+        public TensorFlowTransforms TensorFlow { get; }
+
         internal TransformsCatalog(IHostEnvironment env)
         {
             Contracts.AssertValue(env);
@@ -47,6 +52,7 @@ internal TransformsCatalog(IHostEnvironment env)
             Text = new TextTransforms(this);
             Projection = new ProjectionTransforms(this);
             FeatureSelection = new FeatureSelectionTransforms(this);
+            TensorFlow = new TensorFlowTransforms(this);
         }
 
         public abstract class SubCatalogBase
@@ -109,5 +115,15 @@ internal FeatureSelectionTransforms(TransformsCatalog owner) : base(owner)
             {
             }
         }
+
+        /// <summary>
+        /// The catalog of TensorFlow operations.
+        /// </summary>
+        public sealed class TensorFlowTransforms : SubCatalogBase
+        {
+            internal TensorFlowTransforms(TransformsCatalog owner) : base(owner)
+            {
+            }
+        }
     }
 }

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

@@ -17,7 +17,7 @@ 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 new MLContext().Transforms.TensorFlow.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}");

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,7 +94,7 @@ 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);
@@ -109,7 +109,7 @@ public static DataViewSchema GetModelSchema(IHostEnvironment env, string modelPa
         /// <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)
+        internal static IEnumerable<(string, string, DataViewType, string[])> GetModelNodes(IHostEnvironment env, string modelPath)
         {
             var schema = GetModelSchema(env, modelPath);
 
@@ -338,7 +338,7 @@ 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 TensorFlowModelInfo LoadTensorFlowModel(IHostEnvironment env, string modelPath)
         {
             var session = GetSession(env, modelPath);
             return new TensorFlowModelInfo(env, session, modelPath);

src/Microsoft.ML.TensorFlow/TensorFlowModelInfo.cs

@@ -20,7 +20,7 @@ namespace Microsoft.ML.Transforms
     ///    </item>
     /// </list>
     /// </summary>
-    public class TensorFlowModelInfo
+    public sealed class TensorFlowModelInfo
     {
         internal TFSession Session { get; }
         public string ModelPath { get; }

src/Microsoft.ML.TensorFlow/TensorflowCatalog.cs

@@ -2,9 +2,11 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using System.Collections.Generic;
 using Microsoft.Data.DataView;
 using Microsoft.ML.Data;
 using Microsoft.ML.Transforms;
+using Microsoft.ML.Transforms.TensorFlow;
 
 namespace Microsoft.ML
 {
@@ -25,7 +27,7 @@ public static class TensorflowCatalog
         /// ]]>
         /// </format>
         /// </example>
-        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog catalog,
+        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog.TensorFlowTransforms catalog,
             string modelLocation,
             string outputColumnName,
             string inputColumnName)
@@ -45,7 +47,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// ]]>
         /// </format>
         /// </example>
-        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog catalog,
+        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog.TensorFlowTransforms catalog,
             string modelLocation,
             string[] outputColumnNames,
             string[] inputColumnNames)
@@ -58,7 +60,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// <param name="tensorFlowModel">The pre-loaded TensorFlow model.</param>
         /// <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,
+        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog.TensorFlowTransforms catalog,
             TensorFlowModelInfo tensorFlowModel,
             string outputColumnName,
             string inputColumnName)
@@ -78,7 +80,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// ]]>
         /// </format>
         /// </example>
-        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog catalog,
+        public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog.TensorFlowTransforms catalog,
             TensorFlowModelInfo tensorFlowModel,
             string[] outputColumnNames,
             string[] inputColumnNames)
@@ -90,7 +92,7 @@ public static TensorFlowEstimator ScoreTensorFlowModel(this TransformsCatalog ca
         /// </summary>
         /// <param name="catalog">The transform's catalog.</param>
         /// <param name="options">The <see cref="TensorFlowEstimator.Options"/> specifying the inputs and the settings of the <see cref="TensorFlowEstimator"/>.</param>
-        public static TensorFlowEstimator TensorFlow(this TransformsCatalog catalog,
+        public static TensorFlowEstimator TensorFlow(this TransformsCatalog.TensorFlowTransforms catalog,
             TensorFlowEstimator.Options options)
             => new TensorFlowEstimator(CatalogUtils.GetEnvironment(catalog), options);
 
@@ -100,9 +102,42 @@ public static TensorFlowEstimator TensorFlow(this TransformsCatalog catalog,
         /// <param name="catalog">The transform's catalog.</param>
         /// <param name="options">The <see cref="TensorFlowEstimator.Options"/> specifying the inputs and the settings of the <see cref="TensorFlowEstimator"/>.</param>
         /// <param name="tensorFlowModel">The pre-loaded TensorFlow model.</param>
-        public static TensorFlowEstimator TensorFlow(this TransformsCatalog catalog,
+        public static TensorFlowEstimator TensorFlow(this TransformsCatalog.TensorFlowTransforms catalog,
             TensorFlowEstimator.Options options,
             TensorFlowModelInfo tensorFlowModel)
             => new TensorFlowEstimator(CatalogUtils.GetEnvironment(catalog), options, tensorFlowModel);
+
+        /// <summary>
+        /// This method retrieves the information about the graph nodes of a TensorFlow model as an <see cref="DataViewSchema"/>.
+        /// For every node in the graph that has an output type that is compatible with the types supported by
+        /// <see cref="TensorFlowEstimator"/>, the output schema contains a column with the name of that node, and the
+        /// type of its output (including the item type and the shape, if it is known). Every column also contains metadata
+        /// of kind <see cref="TensorFlowUtils.TensorflowOperatorTypeKind"/>, indicating the operation type of the node, and if that node has inputs in the graph,
+        /// it contains metadata of kind <see cref="TensorFlowUtils.TensorflowUpstreamOperatorsKind"/>, indicating the names of the input nodes.
+        /// </summary>
+        /// <param name="catalog">The transform's catalog.</param>
+        /// <param name="modelLocation">Location of the TensorFlow model.</param>
+        public static DataViewSchema GetModelSchema(this TransformsCatalog.TensorFlowTransforms catalog, string modelLocation)
+            => TensorFlowUtils.GetModelSchema(CatalogUtils.GetEnvironment(catalog), modelLocation);
+
+        /// <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(TransformsCatalog.TensorFlowTransforms, 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="catalog">The transform's catalog.</param>
+        /// <param name="modelLocation">Location of the TensorFlow model.</param>
+        public static IEnumerable<(string, string, DataViewType, string[])> GetModelNodes(this TransformsCatalog.TensorFlowTransforms catalog, string modelLocation)
+            => TensorFlowUtils.GetModelNodes(CatalogUtils.GetEnvironment(catalog), modelLocation);
+
+        /// <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.TensorFlowTransforms, TensorFlowEstimator.Options, TensorFlowModelInfo)"/>.
+        /// </summary>
+        /// <param name="catalog">The transform's catalog.</param>
+        /// <param name="modelLocation">Location of the TensorFlow model.</param>
+        public static TensorFlowModelInfo LoadTensorFlowModel(this TransformsCatalog.TensorFlowTransforms catalog, string modelLocation)
+            => TensorFlowUtils.LoadTensorFlowModel(CatalogUtils.GetEnvironment(catalog), modelLocation);
     }
 }

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

@@ -36,7 +36,7 @@ public void TensorFlowTransforCifarEndToEndTest()
             var pipeEstimator = new ImageLoadingEstimator(mlContext, imageFolder, ("ImageReal", "ImagePath"))
                     .Append(new ImageResizingEstimator(mlContext, "ImageCropped", imageHeight, imageWidth, "ImageReal"))
                     .Append(new ImagePixelExtractingEstimator(mlContext, "Input", "ImageCropped", interleave: true))
-                    .Append(mlContext.Transforms.ScoreTensorFlowModel(model_location, "Output", "Input"))
+                    .Append(mlContext.Transforms.TensorFlow.ScoreTensorFlowModel(model_location, "Output", "Input"))
                     .Append(new ColumnConcatenatingEstimator(mlContext, "Features", "Output"))
                     .Append(new ValueToKeyMappingEstimator(mlContext, "Label"))
                     .AppendCacheCheckpoint(mlContext)