Added support for inserting batch dimension in inputs in TensorFlow. (#2935)

Author: zeahmed

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

@@ -14,14 +14,14 @@ private sealed class OutColumn : Vector<float>
         {
             public PipelineColumn Input { get; }
 
-            public OutColumn(Vector<float> input, string modelFile)
-                : base(new Reconciler(modelFile), input)
+            public OutColumn(Vector<float> input, string modelFile, bool addBatchDimensionInput)
+                : base(new Reconciler(modelFile, addBatchDimensionInput), input)
             {
                 Input = input;
             }
 
-            public OutColumn(Vector<float> input, TensorFlowModel tensorFlowModel)
-                : base(new Reconciler(tensorFlowModel), input)
+            public OutColumn(Vector<float> input, TensorFlowModel tensorFlowModel, bool addBatchDimensionInput)
+                : base(new Reconciler(tensorFlowModel, addBatchDimensionInput), input)
             {
                 Input = input;
             }
@@ -31,20 +31,23 @@ private sealed class Reconciler : EstimatorReconciler
         {
             private readonly string _modelFile;
             private readonly TensorFlowModel _tensorFlowModel;
+            private readonly bool _addBatchDimensionInput;
 
-            public Reconciler(string modelFile)
+            public Reconciler(string modelFile, bool addBatchDimensionInput)
             {
                 Contracts.AssertNonEmpty(modelFile);
                 _modelFile = modelFile;
                 _tensorFlowModel = null;
+                _addBatchDimensionInput = addBatchDimensionInput;
             }
 
-            public Reconciler(TensorFlowModel tensorFlowModel)
+            public Reconciler(TensorFlowModel tensorFlowModel, bool addBatchDimensionInput)
             {
                 Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));
 
                 _modelFile = null;
                 _tensorFlowModel = tensorFlowModel;
+                _addBatchDimensionInput = addBatchDimensionInput;
             }
 
             public override IEstimator<ITransformer> Reconcile(IHostEnvironment env,
@@ -57,9 +60,9 @@ public override IEstimator<ITransformer> Reconcile(IHostEnvironment env,
 
                 var outCol = (OutColumn)toOutput[0];
                 if (_modelFile == null)
-                    return new TensorFlowEstimator(env, new[] { outputNames[outCol] }, new[] { inputNames[outCol.Input] }, _tensorFlowModel);
+                    return new TensorFlowEstimator(env, new[] { outputNames[outCol] }, new[] { inputNames[outCol.Input] }, _tensorFlowModel, _addBatchDimensionInput);
                 else
-                    return new TensorFlowEstimator(env, new[] { outputNames[outCol] }, new[] { inputNames[outCol.Input] }, _modelFile);
+                    return new TensorFlowEstimator(env, new[] { outputNames[outCol] }, new[] { inputNames[outCol.Input] }, _modelFile, _addBatchDimensionInput);
             }
         }
 
@@ -70,22 +73,22 @@ public override IEstimator<ITransformer> Reconcile(IHostEnvironment env,
         /// 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)
+        public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, string modelFile, bool addBatchDimensionInput = false)
         {
             Contracts.CheckValue(input, nameof(input));
             Contracts.CheckNonEmpty(modelFile, nameof(modelFile));
-            return new OutColumn(input, modelFile);
+            return new OutColumn(input, modelFile, addBatchDimensionInput);
         }
 
         /// <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, TensorFlowModel tensorFlowModel)
+        public static Vector<float> ApplyTensorFlowGraph(this Vector<float> input, TensorFlowModel tensorFlowModel, bool addBatchDimensionInput = false)
         {
             Contracts.CheckValue(input, nameof(input));
             Contracts.CheckValue(tensorFlowModel, nameof(tensorFlowModel));
-            return new OutColumn(input, tensorFlowModel);
+            return new OutColumn(input, tensorFlowModel, addBatchDimensionInput);
         }
     }
 }

src/Microsoft.ML.TensorFlow/TensorFlowModel.cs

@@ -55,30 +55,34 @@ public DataViewSchema GetInputSchema()
         /// </summary>
         /// <param name="inputColumnName"> The name of the model input.</param>
         /// <param name="outputColumnName">The name of the requested model output.</param>
+        /// <param name="addBatchDimensionInput">Add a batch dimension to the input e.g. input = [224, 224, 3] => [-1, 224, 224, 3].
+        /// This parameter is used to deal with models that have unknown shape but the internal operators in the model require data to have batch dimension as well.</param>
         /// <example>
         /// <format type="text/markdown">
         /// <![CDATA[
         /// [!code-csharp[ScoreTensorFlowModel](~/../docs/samples/docs/samples/Microsoft.ML.Samples/Dynamic/TensorFlowTransform.cs)]
         /// ]]>
         /// </format>
         /// </example>
-        public TensorFlowEstimator ScoreTensorFlowModel(string outputColumnName, string inputColumnName)
-            => new TensorFlowEstimator(_env, new[] { outputColumnName }, new[] { inputColumnName }, this);
+        public TensorFlowEstimator ScoreTensorFlowModel(string outputColumnName, string inputColumnName, bool addBatchDimensionInput = false)
+            => new TensorFlowEstimator(_env, new[] { outputColumnName }, new[] { inputColumnName }, this, addBatchDimensionInput);
 
         /// <summary>
         /// Scores a dataset using a pre-traiend TensorFlow model.
         /// </summary>
         /// <param name="inputColumnNames"> The names of the model inputs.</param>
         /// <param name="outputColumnNames">The names of the requested model outputs.</param>
+        /// <param name="addBatchDimensionInput">Add a batch dimension to the input e.g. input = [224, 224, 3] => [-1, 224, 224, 3].
+        /// This parameter is used to deal with models that have unknown shape but the internal operators in the model require data to have batch dimension as well.</param>
         /// <example>
         /// <format type="text/markdown">
         /// <![CDATA[
         /// [!code-csharp[ScoreTensorFlowModel](~/../docs/samples/docs/samples/Microsoft.ML.Samples/Dynamic/TensorFlow/ImageClassification.cs)]
         /// ]]>
         /// </format>
         /// </example>
-        public TensorFlowEstimator ScoreTensorFlowModel(string[] outputColumnNames, string[] inputColumnNames)
-            => new TensorFlowEstimator(_env, outputColumnNames, inputColumnNames, this);
+        public TensorFlowEstimator ScoreTensorFlowModel(string[] outputColumnNames, string[] inputColumnNames, bool addBatchDimensionInput = false)
+            => new TensorFlowEstimator(_env, outputColumnNames, inputColumnNames, this, addBatchDimensionInput);
 
         /// <summary>
         /// Retrain the TensorFlow model on new data.
@@ -97,6 +101,8 @@ public TensorFlowEstimator ScoreTensorFlowModel(string[] outputColumnNames, stri
         /// <param name="metricOperation">The name of the operation in the TensorFlow graph to compute performance metric during training (Optional).</param>
         /// <param name="learningRateOperation">The name of the operation in the TensorFlow graph which sets optimizer learning rate (Optional).</param>
         /// <param name="learningRate">Learning rate to use during optimization (Optional).</param>
+        /// <param name="addBatchDimensionInput">Add a batch dimension to the input e.g. input = [224, 224, 3] => [-1, 224, 224, 3].
+        /// This parameter is used to deal with models that have unknown shape but the internal operators in the model require data to have batch dimension as well.</param>
         /// <remarks>
         /// The support for retraining is experimental.
         /// </remarks>
@@ -111,7 +117,8 @@ public TensorFlowEstimator RetrainTensorFlowModel(
             string lossOperation = null,
             string metricOperation = null,
             string learningRateOperation = null,
-            float learningRate = 0.01f)
+            float learningRate = 0.01f,
+            bool addBatchDimensionInput = false)
         {
             var options = new TensorFlowEstimator.Options()
             {
@@ -127,7 +134,8 @@ public TensorFlowEstimator RetrainTensorFlowModel(
                 LearningRateOperation = learningRateOperation,
                 LearningRate = learningRate,
                 BatchSize = batchSize,
-                ReTrain = true
+                ReTrain = true,
+                AddBatchDimensionInputs = addBatchDimensionInput
             };
             return new TensorFlowEstimator(_env, options, this);
         }

src/Microsoft.ML.TensorFlow/TensorflowCatalog.cs

@@ -14,7 +14,7 @@ public static class TensorflowCatalog
     {
         /// <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="TensorFlowModel.ScoreTensorFlowModel(string, string)"/>.
+        /// <see cref="TensorFlowEstimator"/> using <see cref="TensorFlowModel.ScoreTensorFlowModel(string, string, bool)"/>.
         /// </summary>
         /// <param name="catalog">The transform's catalog.</param>
         /// <param name="modelLocation">Location of the TensorFlow model.</param>