Fixing saving/loading with new parameter.

Author: michaelgsharp

src/Microsoft.ML.TensorFlow/TensorflowTransform.cs

@@ -45,6 +45,7 @@ public sealed class TensorFlowTransformer : RowToRowTransformerBase, IDisposable
         private readonly string _savedModelPath;
         private readonly bool _isTemporarySavedModel;
         private readonly bool _addBatchDimensionInput;
+        private readonly bool _treatOutputAsBatched;
         internal readonly Session Session;
         internal readonly Runner Runner;
         internal readonly DataViewType[] OutputTypes;
@@ -71,8 +72,9 @@ private static VersionInfo GetVersionInfo()
                 modelSignature: "TENSFLOW",
                 //verWrittenCur: 0x00010001, // Initial
                 //verWrittenCur: 0x00010002,  // Added Support for Multiple Outputs and SavedModel.
-                verWrittenCur: 0x00010003,  // Added Support for adding batch dimension in inputs.
-                verReadableCur: 0x00010003,
+                //verWrittenCur: 0x00010003,  // Added Support for adding batch dimension in inputs.
+                verWrittenCur: 0x00010004,  // Added Support for treating batch as output or not.
+                verReadableCur: 0x00010004,
                 verWeCanReadBack: 0x00010001,
                 loaderSignature: LoaderSignature,
                 loaderAssemblyName: typeof(TensorFlowTransformer).Assembly.FullName);
@@ -123,20 +125,21 @@ private static TensorFlowTransformer Create(IHostEnvironment env, ModelLoadConte
             // *** Binary format ***
             // byte: indicator for frozen models
             // byte: indicator for adding batch dimension in input
+            // byte: indicator for treating output as batched
             // stream: tensorFlow model.
             // int: number of input columns
             // for each input column
             //   int: id of int column name
             // int: number of output columns
             // for each output column
             //   int: id of output column name
-            GetModelInfo(env, ctx, out string[] inputs, out string[] outputs, out bool isFrozen, out bool addBatchDimensionInput);
+            GetModelInfo(env, ctx, out string[] inputs, out string[] outputs, out bool isFrozen, out bool addBatchDimensionInput, out bool treatOutputAsBatched);
             if (isFrozen)
             {
                 byte[] modelBytes = null;
                 if (!ctx.TryLoadBinaryStream("TFModel", r => modelBytes = r.ReadByteArray()))
                     throw env.ExceptDecode();
-                return new TensorFlowTransformer(env, LoadTFSession(env, modelBytes), outputs, inputs, null, false, addBatchDimensionInput);
+                return new TensorFlowTransformer(env, LoadTFSession(env, modelBytes), outputs, inputs, null, false, addBatchDimensionInput, treatOutputAsBatched: treatOutputAsBatched);
             }
 
             var tempDirPath = Path.GetFullPath(Path.Combine(Path.GetTempPath(), nameof(TensorFlowTransformer) + "_" + Guid.NewGuid()));
@@ -165,7 +168,7 @@ private static TensorFlowTransformer Create(IHostEnvironment env, ModelLoadConte
                     }
                 });
 
-                return new TensorFlowTransformer(env, GetSession(env, tempDirPath), outputs, inputs, tempDirPath, true, addBatchDimensionInput);
+                return new TensorFlowTransformer(env, GetSession(env, tempDirPath), outputs, inputs, tempDirPath, true, addBatchDimensionInput, treatOutputAsBatched: treatOutputAsBatched);
             }
             catch (Exception)
             {
@@ -237,7 +240,7 @@ private static IDataTransform Create(IHostEnvironment env, ModelLoadContext ctx,
         private static IRowMapper Create(IHostEnvironment env, ModelLoadContext ctx, DataViewSchema inputSchema)
             => Create(env, ctx).MakeRowMapper(inputSchema);
 
-        private static void GetModelInfo(IHostEnvironment env, ModelLoadContext ctx, out string[] inputs, out string[] outputs, out bool isFrozen, out bool addBatchDimensionInput)
+        private static void GetModelInfo(IHostEnvironment env, ModelLoadContext ctx, out string[] inputs, out string[] outputs, out bool isFrozen, out bool addBatchDimensionInput, out bool treatOutputAsBatched)
         {
             isFrozen = true;
             bool isNonFrozenModelSupported = ctx.Header.ModelVerReadable >= 0x00010002;
@@ -249,6 +252,11 @@ private static void GetModelInfo(IHostEnvironment env, ModelLoadContext ctx, out
             if (isAddingBatchDimensionSupported)
                 addBatchDimensionInput = ctx.Reader.ReadBoolByte();
 
+            treatOutputAsBatched = true;
+            bool isTreatingOutputAsBatchedSupported = ctx.Header.ModelVerReadable >= 0x00010004;
+            if (isTreatingOutputAsBatchedSupported)
+                treatOutputAsBatched = ctx.Reader.ReadBoolByte();
+
             var numInputs = ctx.Reader.ReadInt32();
             env.CheckDecode(numInputs > 0);
             inputs = new string[numInputs];
@@ -280,6 +288,7 @@ internal TensorFlowTransformer(IHostEnvironment env, Session session, string[] o
             _isTemporarySavedModel = isTemporarySavedModel;
             Session = session;
             _addBatchDimensionInput = addBatchDimensionInput;
+            _treatOutputAsBatched = treatOutputAsBatched;
             Inputs = inputColumnNames;
             Outputs = outputColumnNames;
             tf.compat.v1.disable_eager_execution();
@@ -421,6 +430,7 @@ private protected override void SaveModel(ModelSaveContext ctx)
             // *** Binary format ***
             // byte: indicator for frozen models
             // byte: indicator for adding batch dimension in input
+            // byte: indicator for treating output as batched
             // stream: tensorFlow model.
             // int: number of input columns
             // for each input column
@@ -431,6 +441,7 @@ private protected override void SaveModel(ModelSaveContext ctx)
             var isFrozen = string.IsNullOrEmpty(_savedModelPath);
             ctx.Writer.WriteBoolByte(isFrozen);
             ctx.Writer.WriteBoolByte(_addBatchDimensionInput);
+            ctx.Writer.WriteBoolByte(_treatOutputAsBatched);
             if (isFrozen)
             {
                 using (var status = new Status())

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

@@ -1152,38 +1152,6 @@ public void TensorFlowGettingSchemaMultipleTimes()
             }
         }
 
-        // This test has been created as result of https://github.com/dotnet/machinelearning/issues/5364.
-        [TensorFlowFact]
-        public void TreatOutputAsBatched()
-        {
-            MLContext mlContext = new MLContext();
-
-            var dataView = mlContext.Data.CreateTextLoader<TextInput>().Load(new MultiFileSource(null));
-
-            string modelLocation = @"model_string_test";
-
-            // When treatOutputAsBatched is defaulted to true, make sure the output is correct.
-            using var model = mlContext.Model.LoadTensorFlowModel(modelLocation).ScoreTensorFlowModel(new[] { "Original_A", "Joined_Splited_Text" }, new[] { "A", "B" })
-                .Fit(dataView);
-
-            var modelSchema = model.GetOutputSchema(dataView.Schema);
-
-            Assert.Equal(4, modelSchema.Count);
-            Assert.Equal(new VectorDataViewType(TextDataViewType.Instance, 2), modelSchema[2].Type);
-            Assert.Equal(new VectorDataViewType(TextDataViewType.Instance, 1,1), modelSchema[3].Type);
-
-            using var modelNonBatched = mlContext.Model.LoadTensorFlowModel(modelLocation, false).ScoreTensorFlowModel(new[] { "Original_A", "Joined_Splited_Text" }, new[] { "A", "B" })
-                .Fit(dataView);
-
-            modelSchema = modelNonBatched.GetOutputSchema(dataView.Schema);
-
-            Assert.Equal(4, modelSchema.Count);
-            Assert.Equal(new VectorDataViewType(TextDataViewType.Instance, 0, 2), modelSchema[2].Type);
-            Assert.Equal(new VectorDataViewType(TextDataViewType.Instance, 1, 1), modelSchema[3].Type);
-
-        }
-
-
         [TensorFlowFact]
         public void TensorFlowTransformCifarInvalidShape()
         {

test/Microsoft.ML.Tests/TensorFlowEstimatorTests.cs

@@ -183,6 +183,52 @@ public void TestTensorFlow()
                 Assert.Equal(4, numRows);
             }
         }
+        
+        [TensorFlowFact]
+        public void TreatOutputAsBatched()
+        {
+            var modelLocation = "cifar_model/frozen_model.pb";
+
+            var mlContext = new MLContext(seed: 1);
+            var imageHeight = 32;
+            var imageWidth = 32;
+            var dataFile = GetDataPath("images/images.tsv");
+            var imageFolder = Path.GetDirectoryName(dataFile);
+
+            var data = ML.Data.LoadFromTextFile(dataFile, new[] {
+                new TextLoader.Column("imagePath", DataKind.String, 0),
+                new TextLoader.Column("name", DataKind.String, 1)
+            });
+
+            // Note that CamelCase column names are there to match the TF graph node names.
+            // Check and make sure save/load work correctly for the new TreatOutputAsBatched value.
+            var pipe = ML.Transforms.LoadImages("Input", imageFolder, "imagePath")
+                .Append(ML.Transforms.ResizeImages("Input", imageHeight, imageWidth))
+                .Append(ML.Transforms.ExtractPixels("Input", interleavePixelColors: true))
+                .Append(ML.Model.LoadTensorFlowModel(modelLocation, false).ScoreTensorFlowModel("Output", "Input"));
+
+            TestEstimatorCore(pipe, data);
+            var schema = pipe.Fit(data).Transform(data).Schema;
+
+            // The dimensions of the output with treatOutputAsBatched set to false should be * 10
+            // as the first dimension of -1 is treated as an unkown dimension.
+            Assert.Equal(new VectorDataViewType(NumberDataViewType.Single, 0, 10), schema["Output"].Type);
+
+            // Note that CamelCase column names are there to match the TF graph node names.
+            // Test with TreatOutputAsBatched set to default value of true.
+            pipe = ML.Transforms.LoadImages("Input", imageFolder, "imagePath")
+                .Append(ML.Transforms.ResizeImages("Input", imageHeight, imageWidth))
+                .Append(ML.Transforms.ExtractPixels("Input", interleavePixelColors: true))
+                .Append(ML.Model.LoadTensorFlowModel(modelLocation).ScoreTensorFlowModel("Output", "Input"));
+
+            TestEstimatorCore(pipe, data);
+            schema = pipe.Fit(data).Transform(data).Schema;
+
+            // The dimensions of the output with treatOutputAsBatched set to true should be 10
+            // as the first dimension of -1 is treated as the batch dimension.
+            Assert.Equal(new VectorDataViewType(NumberDataViewType.Single, 10), schema["Output"].Type);
+
+        }
 
         [TensorFlowFact]
         public void TestTensorFlowWithSchema()