From 304957aa8dc8b9742b2e6137dbc1fee6e15fab30 Mon Sep 17 00:00:00 2001 From: Ping He Date: Sun, 16 Jun 2024 18:12:17 -0500 Subject: [PATCH] Added an option to use TensorFlow callback funcs in regModel (#648) * Added call back funcs for regModel. * Enabled ExternalTensorFlow with multiple models. * Fixed a bug for tensorflow. * Fixed an issue for feature ordering. * Fixed a test bug. * Fixed a bug for SSTFIML. * Allowed writeFeatures for FI mode. * Added the nInputs back. --- dafoam/pyDAFoam.py | 53 ++++-- .../DATurbulenceModel/DAkOmegaSSTFIML.C | 4 + src/adjoint/DARegression/DARegression.C | 156 +++++++++++++++++- src/adjoint/DARegression/DARegression.H | 55 ++++++ src/adjoint/DASolver/DASolver.C | 3 + src/adjoint/DASolver/DASolver.H | 7 +- src/adjoint/DAUtility/DAUtility.H | 4 + src/pyDASolvers/DASolvers.H | 7 +- src/pyDASolvers/pyDASolvers.pyx | 10 +- tests/runTests_DASimpleFoamkOmegaSSTFIML.py | 7 +- 10 files changed, 270 insertions(+), 36 deletions(-) diff --git a/dafoam/pyDAFoam.py b/dafoam/pyDAFoam.py index f8e74910..484e0df8 100755 --- a/dafoam/pyDAFoam.py +++ b/dafoam/pyDAFoam.py @@ -743,10 +743,9 @@ def __init__(self): ## tensorflow related functions self.tensorflow = { "active": False, - "modelName": "model", - "nInputs": 1, - "nOutputs": 1, - "batchSize": 1000, + #"model1": { + # "predictBatchSize": 1000 + #} } @@ -928,9 +927,9 @@ def __init__(self, comm=None, options=None): TensorFlowHelper.options = self.getOption("tensorflow") TensorFlowHelper.initialize() # pass this helper function to the C++ layer - self.solver.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd) + self.solver.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName) if self.getOption("useAD")["mode"] in ["forward", "reverse"]: - self.solverAD.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd) + self.solverAD.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName) Info("pyDAFoam initialization done!") @@ -4491,7 +4490,13 @@ class TensorFlowHelper: options = {} - model = None + model = {} + + modelName = None + + predictBatchSize = {} + + nInputs = {} @staticmethod def initialize(): @@ -4499,14 +4504,19 @@ def initialize(): Initialize parameters and load models """ Info("Initializing the TensorFlowHelper") - TensorFlowHelper.modelName = TensorFlowHelper.options["modelName"] - TensorFlowHelper.nInputs = TensorFlowHelper.options["nInputs"] - TensorFlowHelper.nOutputs = TensorFlowHelper.options["nOutputs"] - TensorFlowHelper.batchSize = TensorFlowHelper.options["batchSize"] - TensorFlowHelper.model = tf.keras.models.load_model(TensorFlowHelper.modelName) + for key in list(TensorFlowHelper.options.keys()): + if key != "active": + modelName = key + TensorFlowHelper.predictBatchSize[modelName] = TensorFlowHelper.options[modelName]["predictBatchSize"] + TensorFlowHelper.nInputs[modelName] = TensorFlowHelper.options[modelName]["nInputs"] + TensorFlowHelper.model[modelName] = tf.keras.models.load_model(modelName) - if TensorFlowHelper.nOutputs != 1: - raise Error("current version supports nOutputs=1 only!") + @staticmethod + def setModelName(modelName): + """ + Set the model name from the C++ to Python layer + """ + TensorFlowHelper.modelName = modelName.decode() @staticmethod def predict(inputs, n, outputs, m): @@ -4514,8 +4524,12 @@ def predict(inputs, n, outputs, m): Calculate the outputs based on the inputs using the saved model """ - inputs_tf = np.reshape(inputs, (-1, TensorFlowHelper.nInputs)) - outputs_tf = TensorFlowHelper.model.predict(inputs_tf, verbose=False, batch_size=TensorFlowHelper.batchSize) + modelName = TensorFlowHelper.modelName + nInputs = TensorFlowHelper.nInputs[modelName] + + inputs_tf = np.reshape(inputs, (-1, nInputs)) + batchSize = TensorFlowHelper.predictBatchSize[modelName] + outputs_tf = TensorFlowHelper.model[modelName].predict(inputs_tf, verbose=False, batch_size=batchSize) for i in range(m): outputs[i] = outputs_tf[i, 0] @@ -4526,11 +4540,14 @@ def calcJacVecProd(inputs, inputs_b, n, outputs, outputs_b, m): Calculate the gradients of the outputs wrt the inputs """ - inputs_tf = np.reshape(inputs, (-1, TensorFlowHelper.nInputs)) + modelName = TensorFlowHelper.modelName + nInputs = TensorFlowHelper.nInputs[modelName] + + inputs_tf = np.reshape(inputs, (-1, nInputs)) inputs_tf_var = tf.Variable(inputs_tf, dtype=tf.float32) with tf.GradientTape() as tape: - outputs_tf = TensorFlowHelper.model(inputs_tf_var) + outputs_tf = TensorFlowHelper.model[modelName](inputs_tf_var) gradients_tf = tape.gradient(outputs_tf, inputs_tf_var) diff --git a/src/adjoint/DAModel/DATurbulenceModel/DAkOmegaSSTFIML.C b/src/adjoint/DAModel/DATurbulenceModel/DAkOmegaSSTFIML.C index dc47e2d2..a0de7cc0 100755 --- a/src/adjoint/DAModel/DATurbulenceModel/DAkOmegaSSTFIML.C +++ b/src/adjoint/DAModel/DATurbulenceModel/DAkOmegaSSTFIML.C @@ -885,6 +885,10 @@ void DAkOmegaSSTFIML::calcBetaField() inputs_[cI * 9 + 8] = tauRatio_[cI]; } + // set the model name in the Python layer + word modelName = "model"; + DAUtility::pySetModelNameInterface(modelName.c_str(), DAUtility::pySetModelName); + // NOTE: forward mode not supported.. #if defined(CODI_AD_REVERSE) diff --git a/src/adjoint/DARegression/DARegression.C b/src/adjoint/DARegression/DARegression.C index 5ed93a2c..23291beb 100755 --- a/src/adjoint/DARegression/DARegression.C +++ b/src/adjoint/DARegression/DARegression.C @@ -87,9 +87,19 @@ DARegression::DARegression( { nRBFs_.set(modelName, modelSubDict.getLabel("nRBFs")); } + else if (modelType_[modelName] == "externalTensorFlow") + { + useExternalModel_ = 1; + // here the array size is chosen based on the regModel that has the largest number of inputs + // this is important because we support multiple regModels but we don't want to create multiple featuresFlattenArray_ + if (inputNames_[modelName].size() * mesh_.nCells() > featuresFlattenArraySize_) + { + featuresFlattenArraySize_ = inputNames_[modelName].size() * mesh_.nCells(); + } + } else { - FatalErrorIn("DARegression") << "modelType_: " << modelType_[modelName] << " not supported. Options are: neuralNetwork and radialBasisFunction" << abort(FatalError); + FatalErrorIn("DARegression") << "modelType_: " << modelType_[modelName] << " not supported. Options are: neuralNetwork, radialBasisFunction, and externalTensorFlow" << abort(FatalError); } // check the sizes @@ -123,6 +133,19 @@ DARegression::DARegression( } features_.set(modelName, features); } + + // if external model is used, initialize space for the features and output arrays + if (useExternalModel_) + { + +#if defined(CODI_AD_FORWARD) + featuresFlattenArrayDouble_ = new double[featuresFlattenArraySize_]; + outputFieldArrayDouble_ = new double[mesh_.nCells()]; +#else + featuresFlattenArray_ = new scalar[featuresFlattenArraySize_]; + outputFieldArray_ = new scalar[mesh_.nCells()]; +#endif + } } } @@ -341,15 +364,15 @@ label DARegression::compute() { word modelName = modelNames_[idxI]; + // compute the inputFeature for all inputs + this->calcInputFeatures(modelName); + // if the output variable is not found in the Db, just return and do nothing if (!mesh_.thisDb().foundObject(outputName_[modelName])) { return 0; } - // compute the inputFeature for all inputs - this->calcInputFeatures(modelName); - volScalarField& outputField = const_cast(mesh_.thisDb().lookupObject(outputName_[modelName])); if (modelType_[modelName] == "neuralNetwork") @@ -472,9 +495,126 @@ label DARegression::compute() outputField.correctBoundaryConditions(); } + else if (modelType_[modelName] == "externalTensorFlow") + { + label nInputs = inputNames_[modelName].size(); + label nCells = mesh_.nCells(); + + DAUtility::pySetModelNameInterface(modelName.c_str(), DAUtility::pySetModelName); + + // NOTE: forward mode not supported.. +#if defined(CODI_AD_REVERSE) + + // assign features_ to featuresFlattenArray_ + // here featuresFlattenArray_ should be order like this to facilitate Python layer reshape: + // [(cell1, feature1), (cell1, feature2), ... (cell2, feature1), (cell2, feature2) ... ] + label counterI = 0; + // loop over all features + forAll(features_[modelName], idxI) + { + // loop over all cells + forAll(features_[modelName][idxI], cellI) + { + counterI = cellI * nCells + idxI; + featuresFlattenArray_[counterI] = features_[modelName][idxI][cellI]; + } + } + // assign outputField to outputFieldArray_ + forAll(outputField, cellI) + { + outputFieldArray_[cellI] = outputField[cellI]; + } + + // we need to use the external function helper from CoDiPack to propagate the AD + codi::ExternalFunctionHelper externalFunc; + for (label i = 0; i < mesh_.nCells() * nInputs; i++) + { + externalFunc.addInput(featuresFlattenArray_[i]); + } + + for (label i = 0; i < mesh_.nCells(); i++) + { + externalFunc.addOutput(outputFieldArray_[i]); + } + + externalFunc.callPrimalFunc(DARegression::betaCompute); + + codi::RealReverse::Tape& tape = codi::RealReverse::getTape(); + + if (tape.isActive()) + { + externalFunc.addToTape(DARegression::betaJacVecProd); + } + + forAll(outputField, cellI) + { + outputField[cellI] = outputFieldArray_[cellI]; + } + +#elif defined(CODI_AD_FORWARD) + // assign features_ to featuresFlattenArray_ + // here featuresFlattenArray_ should be order like this to facilitate Python layer reshape: + // [(cell1, feature1), (cell1, feature2), ... (cell2, feature1), (cell2, feature2) ... ] + label counterI = 0; + // loop over all features + forAll(features_[modelName], idxI) + { + // loop over all cells + forAll(features_[modelName][idxI], cellI) + { + counterI = cellI * nCells + idxI; + featuresFlattenArrayDouble_[counterI] = features_[modelName][idxI][cellI].getValue(); + } + } + // assign outputField to outputFieldArrayDouble_ + forAll(outputField, cellI) + { + outputFieldArrayDouble_[cellI] = outputField[cellI].value(); + } + + // python callback function + DAUtility::pyCalcBetaInterface( + featuresFlattenArrayDouble_, mesh_.nCells() * nInputs, outputFieldArrayDouble_, mesh_.nCells(), DAUtility::pyCalcBeta); + + forAll(outputField, cellI) + { + outputField[cellI] = outputFieldArrayDouble_[cellI]; + } + +#else + // assign features_ to featuresFlattenArray_ + // here featuresFlattenArray_ should be order like this to facilitate Python layer reshape: + // [(cell1, feature1), (cell1, feature2), ... (cell2, feature1), (cell2, feature2) ... ] + label counterI = 0; + // loop over all features + forAll(features_[modelName], idxI) + { + // loop over all cells + forAll(features_[modelName][idxI], cellI) + { + counterI = cellI * nCells + idxI; + featuresFlattenArray_[counterI] = features_[modelName][idxI][cellI]; + } + } + // assign outputField to outputFieldArray_ + forAll(outputField, cellI) + { + outputFieldArray_[cellI] = outputField[cellI]; + } + + // python callback function + DAUtility::pyCalcBetaInterface( + featuresFlattenArray_, mesh_.nCells() * nInputs, outputFieldArray_, mesh_.nCells(), DAUtility::pyCalcBeta); + + forAll(outputField, cellI) + { + outputField[cellI] = outputFieldArray_[cellI]; + } +#endif + } else { - FatalErrorIn("") << "modelType_: " << modelType_ << " not supported. Options are: neuralNetwork and radialBasisFunction" << abort(FatalError); + FatalErrorIn("") << "modelType_: " << modelType_ << " not supported. Options are: neuralNetwork, radialBasisFunction, and externalTensorFlow" << abort(FatalError); } } @@ -529,9 +669,13 @@ label DARegression::nParameters(word modelName) return nParameters; } + else if (modelType_[modelName] == "externalTensorFlow") + { + // do nothing + } else { - FatalErrorIn("") << "modelType_: " << modelType_[modelName] << " not supported. Options are: neuralNetwork and radialBasisFunction" << abort(FatalError); + FatalErrorIn("") << "modelType_: " << modelType_[modelName] << " not supported. Options are: neuralNetwork, radialBasisFunction, and externalTensorFlow" << abort(FatalError); } } diff --git a/src/adjoint/DARegression/DARegression.H b/src/adjoint/DARegression/DARegression.H index 7144daa0..77eefed9 100755 --- a/src/adjoint/DARegression/DARegression.H +++ b/src/adjoint/DARegression/DARegression.H @@ -105,6 +105,25 @@ protected: /// a list of scalarField to save the input features HashTable> features_; +#if defined(CODI_AD_FORWARD) + /// a flatten double feature array for externalTensorFlow model (forward AD) + double* featuresFlattenArrayDouble_; + + /// output field double array for externalTensorFlow model (forward AD) + double* outputFieldArrayDouble_; +#else + /// a flatten feature array for externalTensorFlow model + scalar* featuresFlattenArray_; + + /// output field array for externalTensorFlow model + scalar* outputFieldArray_; +#endif + /// whether to use external model + label useExternalModel_ = 0; + + /// the array size is chosen based on the regModel that has the largest number of inputs (this is important because we support multiple regModels but we don't want to create multiple featuresFlattenArray_) + label featuresFlattenArraySize_ = -100; + /// whether to write the feature fields to the disk HashTable