fix: Added support for training a multi-input model using a dataset.

SciSharp · Jul 1, 2024 · 93dda17 · 93dda17
1 parent f8b7bde
commit 93dda17
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs b/src/TensorFlowNET.Keras/Engine/Model.Evaluate.cs
@@ -112,7 +112,19 @@ public Dictionary<string, float> evaluate(IDatasetV2 x, int verbose = 1, bool is
                 Steps = data_handler.Inferredsteps
             });
 
-            return evaluate(data_handler, callbacks, is_val, test_function);
+            Func<DataHandler, OwnedIterator, Dictionary<string, float>> testFunction;
+
+            if (data_handler.DataAdapter.GetDataset().structure.Length > 2 ||
+                data_handler.DataAdapter.GetDataset().FirstInputTensorCount > 1)
+            {
+                testFunction = test_step_multi_inputs_function;
+            }
+            else
+            {
+                testFunction = test_function;
+            }
+
+            return evaluate(data_handler, callbacks, is_val, testFunction);
         }
 
         /// <summary>

diff --git a/src/TensorFlowNET.Keras/Engine/Model.Fit.cs b/src/TensorFlowNET.Keras/Engine/Model.Fit.cs
@@ -179,9 +179,20 @@ public ICallback fit(IDatasetV2 dataset,
                 StepsPerExecution = _steps_per_execution
             });
 
+            Func<DataHandler, OwnedIterator, Dictionary<string, float>> trainStepFunction;
+
+            if (data_handler.DataAdapter.GetDataset().structure.Length > 2 ||
+                data_handler.DataAdapter.GetDataset().FirstInputTensorCount > 1)
+            {
+                trainStepFunction = train_step_multi_inputs_function;
+            }
+            else
+            {
+                trainStepFunction = train_step_function;
+            }
 
             return FitInternal(data_handler, epochs, validation_step, verbose, callbacks, validation_data: validation_data,
-                    train_step_func: train_step_function);
+                    train_step_func: trainStepFunction);
         }
 
         History FitInternal(DataHandler data_handler, int epochs, int validation_step, int verbose, List<ICallback> callbackList, IDatasetV2 validation_data,