Skip to content

Commit

Permalink
Added an option to use TensorFlow callback funcs in regModel (#648)
Browse files Browse the repository at this point in the history 
* 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.
  • Loading branch information
@friedenhe
friedenhe authored Jun 16, 2024
1 parent eb6d1b3 commit 304957a
Show file tree
Hide file tree
Showing 10 changed files with 270 additions and 36 deletions.
53 changes: 35 additions & 18 deletions dafoam/pyDAFoam.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
#}
}


Expand Down Expand Up @@ -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!")

Expand Down Expand Up @@ -4491,31 +4490,46 @@ class TensorFlowHelper:

options = {}

model = None
model = {}

modelName = None

predictBatchSize = {}

nInputs = {}

@staticmethod
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):
"""
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]
Expand All @@ -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)

Expand Down
4 changes: 4 additions & 0 deletions src/adjoint/DAModel/DATurbulenceModel/DAkOmegaSSTFIML.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
156 changes: 150 additions & 6 deletions src/adjoint/DARegression/DARegression.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
}
}

Expand Down Expand Up @@ -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<volScalarField>(outputName_[modelName]))
{
return 0;
}

// compute the inputFeature for all inputs
this->calcInputFeatures(modelName);

volScalarField& outputField = const_cast<volScalarField&>(mesh_.thisDb().lookupObject<volScalarField>(outputName_[modelName]));

if (modelType_[modelName] == "neuralNetwork")
Expand Down Expand Up @@ -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<codi::RealReverse> 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);
}
}

Expand Down Expand Up @@ -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);
}
}

Expand Down
55 changes: 55 additions & 0 deletions src/adjoint/DARegression/DARegression.H
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,25 @@ protected:
/// a list of scalarField to save the input features
HashTable<PtrList<volScalarField>> 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<label> writeFeatures_;

Expand All @@ -118,6 +137,16 @@ public:
/// Destructor
virtual ~DARegression()
{
if (useExternalModel_)
{
#if defined(CODI_AD_FORWARD)
delete[] featuresFlattenArrayDouble_;
delete[] outputFieldArrayDouble_;
#else
delete[] featuresFlattenArray_;
delete[] outputFieldArray_;
#endif
}
}

// Members
Expand Down Expand Up @@ -188,6 +217,32 @@ public:
}
}
}

#ifdef CODI_AD_REVERSE

/// these two functions are for AD external functions
static void betaCompute(
const double* x,
size_t n,
double* y,
size_t m,
codi::ExternalFunctionUserData* d)
{
DAUtility::pyCalcBetaInterface(x, n, y, m, DAUtility::pyCalcBeta);
}

static void betaJacVecProd(
const double* x,
double* x_b,
size_t n,
const double* y,
const double* y_b,
size_t m,
codi::ExternalFunctionUserData* d)
{
DAUtility::pyCalcBetaJacVecProdInterface(x, x_b, n, y, y_b, m, DAUtility::pyCalcBetaJacVecProd);
}
#endif
};

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
Expand Down
3 changes: 3 additions & 0 deletions src/adjoint/DASolver/DASolver.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ pyComputeInterface Foam::DAUtility::pyCalcBetaInterface = NULL;
void* Foam::DAUtility::pyCalcBetaJacVecProd = NULL;
pyJacVecProdInterface Foam::DAUtility::pyCalcBetaJacVecProdInterface = NULL;

void* Foam::DAUtility::pySetModelName = NULL;
pySetCharInterface Foam::DAUtility::pySetModelNameInterface = NULL;

scalar Foam::DAUtility::primalMaxInitRes_ = -1e16;

namespace Foam
Expand Down
Loading

0 comments on commit 304957a

Please sign in to comment.