max/poet
1
0
Fork 0
mirror of https://git.gfz-potsdam.de/naaice/poet.git synced 2025-12-16 12:54:50 +01:00
Code Issues Packages Projects Releases Wiki Activity

Eigen model works

Browse Source
This commit is contained in:
straile 2024-10-15 11:17:30 +02:00
parent e99a114bc3
commit 8691370abb
4 changed files with 205 additions and 165 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py
View File

@ -17,6 +17,9 @@ def training_step(model, x, y, x_val, y_val, batch_size, epochs):
def prediction_step(model, x, batch_size):
prediction = model.predict(x, batch_size)
return np.array(prediction, dtype=np.float64)
#with tf.device('/cpu:0'):
# prediction = model.predict(x, batch_size)
# return np.array(prediction, dtype=np.float64)
def get_weights(model):
weights = model.get_weights()

242
src/Chemistry/SurrogateModels/AI_functions.cpp
View File

@ -36,49 +36,6 @@ int Python_Keras_setup(std::string functions_file_path) {
return py_functions_initialized;
}
/**
* @brief Sets the ai surrogate runtime parameters
* @param R Global R environment
* @param params Gobal runtime parameters struct
* @param init_list List with initial data
*/
void set_ai_surrogate_runtime_params(RInsidePOET& R, RuntimeParameters& params, InitialList& init_list) {
/* AI surrogate training and inference parameters. (Can be set by declaring a
variable of the same name in one of the the R input scripts)*/
params.Keras_training_always_use_CPU = false; // Default will use GPU if detected
params.Keras_training_always_use_CPU = false; // Default will use GPU if detected
params.use_Keras_predictions = false; // Default inference function is custom C++ / Eigen implementation
params.batch_size = 2560; // default value determined in test on the UP Turing cluster
params.training_epochs = 20; //
params.training_data_size = init_list.getDiffusionInit().n_rows *
init_list.getDiffusionInit().n_cols; // Default value is number of cells in field
params.save_model_path = ""; // Model is only saved if a path is set in the input field
if (Rcpp::as<bool>(R.parseEval("exists(\"batch_size\")"))) {
params.batch_size = R["batch_size"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"training_epochs\")"))) {
params.training_epochs = R["training_epochs"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"training_data_size\")"))) {
params.training_data_size = R["training_data_size"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"use_Keras_predictions\")"))) {
params.use_Keras_predictions = R["use_Keras_predictions"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"Keras_predictions_always_use_CPU\")"))) {
params.Keras_predictions_always_use_CPU = R["Keras_predictions_always_use_CPU"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"Keras_training_always_use_CPU\")"))) {
params.Keras_training_always_use_CPU = R["Keras_training_always_use_CPU"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"save_model_path\")"))) {
params.save_model_path = Rcpp::as<std::string>(R["save_model_path"]);
std::cout << "AI: Model will be saved as \"" << params.save_model_path << "\"" << std::endl;
}
}
/**
* @brief Loads the user-supplied Keras model
* @param model_file_path Path to a .keras file that the user must supply as
@ -195,6 +152,28 @@ std::vector<double> Python_Keras_predict(std::vector<std::vector<double>> x, int
return predictions;
}
/**
* @brief Uses Eigen for fast inference with the weights and biases of a neural network
* @param input_batch Batch of input data that must fit the size of the neural networks input layer
* @param model Struct of aligned Eigen vectors that hold the neural networks weights and biases.
* Only supports simple fully connected feed forward networks.
* @return The batch of predictions made with the neural network weights and biases and the data
* in input_batch
*/
Eigen::MatrixXd eigen_inference_batched(const Eigen::Ref<const Eigen::MatrixXd>& input_batch, const EigenModel& model) {
Eigen::MatrixXd current_layer = input_batch;
// Process all hidden layers
for (size_t layer = 0; layer < model.weight_matrices.size() - 1; ++layer) {
current_layer = (model.weight_matrices[layer] * current_layer);
current_layer = current_layer.colwise() + model.biases[layer];
current_layer = current_layer.array().max(0.0);
}
// Process output layer (without ReLU)
size_t output_layer = model.weight_matrices.size() - 1;
return (model.weight_matrices[output_layer] * current_layer).colwise() + model.biases[output_layer];
}
/**
* @brief Uses the Eigen representation of the Keras model weights for fast inference
@ -203,13 +182,12 @@ std::vector<double> Python_Keras_predict(std::vector<std::vector<double>> x, int
* @return Predictions that the neural network made from the input values x. The predictions are
* represented as a vector similar to the representation from the Field.AsVector() method
*/
std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vector<double>> x, int batch_size,
std::vector<double> Eigen_predict(const EigenModel& model, const std::vector<std::vector<double>>& x, int batch_size,
std::mutex* Eigen_model_mutex) {
// Convert input data to Eigen matrix
const int num_samples = x[0].size();
const int num_features = x.size();
Eigen::MatrixXd full_input_matrix(num_features, num_samples);
for (int i = 0; i < num_samples; ++i) {
for (int j = 0; j < num_features; ++j) {
full_input_matrix(j, i) = x[j][i];
@ -217,32 +195,38 @@ std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vect
}
std::vector<double> result;
result.reserve(num_samples * num_features);
result.reserve(num_samples);
if (num_features != model.weight_matrices[0].cols()) {
throw std::runtime_error("Input data size " + std::to_string(num_features) + \
" does not match model input layer of size " + std::to_string(model.weight_matrices[0].cols()));
throw std::runtime_error("Input data size " + std::to_string(num_features) +
" does not match model input layer of size " +
std::to_string(model.weight_matrices[0].cols()));
}
int num_batches = std::ceil(static_cast<double>(num_samples) / batch_size);
Eigen_model_mutex->lock();
std::lock_guard<std::mutex> lock(Eigen_model_mutex);
for (int batch = 0; batch < num_batches; ++batch) {
int start_idx = batch * batch_size;
int end_idx = std::min((batch + 1) * batch_size, num_samples);
int current_batch_size = end_idx - start_idx;
// Extract the current input data batch
Eigen::MatrixXd batch_data(num_features, current_batch_size);
batch_data = full_input_matrix.block(0, start_idx, num_features, current_batch_size);
// Predict
batch_data = eigen_inference_batched(batch_data, model);
result.insert(result.end(), batch_data.data(), batch_data.data() + batch_data.size());
// Extract the current input data batch
Eigen::MatrixXd batch_data = full_input_matrix.block(0, start_idx, num_features, current_batch_size);
// Predict
Eigen::MatrixXd output = eigen_inference_batched(batch_data, model);
// Append the results
result.insert(result.end(), output.data(), output.data() + output.size());
}
Eigen_model_mutex->unlock();
return result;
}
/**
* @brief Appends data from one matrix (column major std::vector<std::vector<double>>) to another
* @param training_data_buffer Matrix that the values are appended to
@ -373,7 +357,7 @@ void parallel_training(EigenModel* Eigen_model,
if (!params.use_Keras_predictions) {
std::cout << "AI: Training thread: Update shared model weights" << std::endl;
Eigen_model_mutex->lock();
Python_Keras_set_weights_as_Eigen(*Eigen_model);
//Python_Keras_set_weights_as_Eigen(Eigen_model);
Eigen_model_mutex->unlock();
}
@ -415,33 +399,32 @@ int Python_Keras_training_thread(EigenModel* Eigen_model,
return 0;
}
/**
* @brief Uses Eigen for fast inference with the weights and biases of a neural network
* @param input_batch Batch of input data that must fit the size of the neural networks input layer
* @param model Struct of aligned Eigen vectors that hold the neural networks weights and biases.
* Only supports simple fully connected feed forward networks.
* @return The batch of predictions made with the neural network weights and biases and the data
* in input_batch
*/
Eigen::MatrixXd eigen_inference_batched(const Eigen::Ref<Eigen::MatrixXd>& input_batch, const EigenModel& model) {
Eigen::MatrixXd current_layer = input_batch;
// Process all hidden layers
for (size_t layer = 0; layer < model.weight_matrices.size() - 1; ++layer) {
std::cout << "LAYER " << layer << std::endl;
current_layer = (model.weight_matrices[layer] * current_layer);
current_layer = current_layer.colwise() + model.biases[layer];
current_layer = current_layer.array().max(0.0);
void update_weights(EigenModel* model,
const std::vector<std::vector<std::vector<double>>>& weights) {
size_t num_layers = weights.size() / 2;
for (size_t i = 0; i < weights.size(); i += 2) {
// Fill current weight matrix
size_t rows = weights[i][0].size();
size_t cols = weights[i].size();
for (size_t j = 0; j < cols; ++j) {
for (size_t k = 0; k < rows; ++k) {
model->weight_matrices[i / 2](k, j) = weights[i][j][k];
}
// Process output layer (without ReLU)
size_t output_layer = model.weight_matrices.size() - 1;
return (model.weight_matrices[output_layer] * current_layer).colwise() + model.biases[output_layer];
}
// Fill bias vector
size_t bias_size = weights[i + 1][0].size();
for (size_t j = 0; j < bias_size; ++j) {
model->biases[i / 2](j) = weights[i + 1][0][j];
}
}
}
/**
* @brief Converts the weights and biases from the Python Keras model to Eigen matrices
* @return A EigenModel struct containing the model weights and biases as aligned Eigen matrices
*/
void Python_Keras_set_weights_as_Eigen(EigenModel& eigen_model) {
std::vector<std::vector<std::vector<double>>> Python_Keras_get_weights() {
// Acquire the Python GIL
PyGILState_STATE gstate = PyGILState_Ensure();
@ -459,72 +442,83 @@ void Python_Keras_set_weights_as_Eigen(EigenModel& eigen_model) {
throw std::runtime_error("Failed to get weights from Keras model");
}
// Clear old values
eigen_model.weight_matrices.clear();
eigen_model.biases.clear();
// Container for the extracted weights
std::vector<std::vector<std::vector<double>>> cpp_weights;
// Iterate through the layers (weights and biases)
Py_ssize_t num_layers = PyList_Size(py_weights_list);
for (Py_ssize_t i = 0; i < num_layers; i += 2) {
// Get weight matrix
PyObject* weight_array = PyList_GetItem(py_weights_list, i);
if (!weight_array) throw std::runtime_error("Failed to get weight array");
if (!PyArray_Check(weight_array)) throw std::runtime_error("Weight array is not a NumPy array");
PyArrayObject* weight_np = reinterpret_cast<PyArrayObject*>(weight_array);
if (PyArray_NDIM(weight_np) != 2) throw std::runtime_error("Weight array is not 2-dimensional");
// Check data type
int dtype = PyArray_TYPE(weight_np);
if (dtype != NPY_FLOAT32 && dtype != NPY_DOUBLE) {
throw std::runtime_error("Unsupported data type.\nMust be NPY_FLOAT32 or NPY_DOUBLE");
for (Py_ssize_t i = 0; i < num_layers; ++i) {
PyObject* py_weight_array = PyList_GetItem(py_weights_list, i);
if (!PyArray_Check(py_weight_array)) {
throw std::runtime_error("Weight is not a NumPy array.");
}
// Cast the data correctly depending on the type
void* weight_data = PyArray_DATA(weight_np);
PyArrayObject* weight_np = reinterpret_cast<PyArrayObject*>(py_weight_array);
int dtype = PyArray_TYPE(weight_np);
if (PyArray_NDIM(weight_np) == 2) {
// Handle 2D weight matrices (e.g., Dense layer weights)
int num_rows = PyArray_DIM(weight_np, 0);
int num_cols = PyArray_DIM(weight_np, 1);
Eigen::MatrixXd weight_matrix;
// Create a container for the matrix
std::vector<std::vector<double>> weight_matrix(num_rows, std::vector<double>(num_cols));
// Handle different data types
if (dtype == NPY_FLOAT32) {
// Handle float32 array from NumPy
float* weight_data_float = static_cast<float*>(weight_data);
weight_matrix = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
weight_data_float, num_rows, num_cols).cast<double>().transpose();
float* weight_data_float = static_cast<float*>(PyArray_DATA(weight_np));
for (int r = 0; r < num_rows; ++r) {
for (int c = 0; c < num_cols; ++c) {
weight_matrix[r][c] = static_cast<double>(weight_data_float[r * num_cols + c]);
}
}
} else if (dtype == NPY_DOUBLE) {
// Handle double (float64) array from NumPy
double* weight_data_double = static_cast<double*>(weight_data);
weight_matrix = Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(
weight_data_double, num_rows, num_cols).transpose();
double* weight_data_double = static_cast<double*>(PyArray_DATA(weight_np));
for (int r = 0; r < num_rows; ++r) {
for (int c = 0; c < num_cols; ++c) {
weight_matrix[r][c] = weight_data_double[r * num_cols + c];
}
}
} else {
throw std::runtime_error("Unsupported data type for weights. Must be NPY_FLOAT32 or NPY_DOUBLE.");
}
// Get bias vector
PyObject* bias_array = PyList_GetItem(py_weights_list, i + 1);
PyArrayObject* bias_np = reinterpret_cast<PyArrayObject*>(bias_array);
if (PyArray_NDIM(bias_np) != 1) throw std::runtime_error("Bias array is not 1-dimensional");
cpp_weights.push_back(weight_matrix);
// Check bias data type and cast accordingly
void* bias_data = PyArray_DATA(bias_np);
int bias_dtype = PyArray_TYPE(bias_np);
int bias_size = PyArray_DIM(bias_np, 0);
Eigen::VectorXd bias_vector;
} else if (PyArray_NDIM(weight_np) == 1) {
// Handle 1D bias vectors
int num_elements = PyArray_DIM(weight_np, 0);
if (bias_dtype == NPY_FLOAT32) {
float* bias_data_float = static_cast<float*>(bias_data);
bias_vector = Eigen::Map<Eigen::Matrix<float, Eigen::Dynamic, 1>>(bias_data_float, bias_size).cast<double>();
} else if (bias_dtype == NPY_DOUBLE) {
double* bias_data_double = static_cast<double*>(bias_data);
bias_vector = Eigen::Map<Eigen::VectorXd>(bias_data_double, bias_size);
// Create a container for the vector
std::vector<std::vector<double>> bias_vector(1, std::vector<double>(num_elements));
// Handle different data types
if (dtype == NPY_FLOAT32) {
float* bias_data_float = static_cast<float*>(PyArray_DATA(weight_np));
for (int j = 0; j < num_elements; ++j) {
bias_vector[0][j] = static_cast<double>(bias_data_float[j]);
}
} else if (dtype == NPY_DOUBLE) {
double* bias_data_double = static_cast<double*>(PyArray_DATA(weight_np));
for (int j = 0; j < num_elements; ++j) {
bias_vector[0][j] = bias_data_double[j];
}
} else {
throw std::runtime_error("Unsupported data type for biases. Must be NPY_FLOAT32 or NPY_DOUBLE.");
}
// Add to EigenModel
eigen_model.weight_matrices.push_back(weight_matrix);
eigen_model.biases.push_back(bias_vector);
cpp_weights.push_back(bias_vector);
} else {
throw std::runtime_error("Unsupported weight dimension. Only 1D and 2D arrays are supported.");
}
}
// Clean up
Py_DECREF(py_weights_list);
Py_DECREF(args);
// Release the Python GIL
// Release Python GIL
PyGILState_Release(gstate);
return cpp_weights;
}

22
src/Chemistry/SurrogateModels/AI_functions.hpp
View File

@ -17,8 +17,6 @@
#include <string>
#include <vector>
#include "poet.hpp"
#include "Base/RInsidePOET.hpp"
#include "Init/InitialList.hpp"
// PhreeqC definition of pi clashes with Eigen macros so we have to temporarily undef it
#pragma push_macro("pi")
@ -27,13 +25,10 @@
#pragma pop_macro("pi")
namespace poet {
// Define an aligned allocator for std::vector
template<typename T>
using aligned_vector = std::vector<T, Eigen::aligned_allocator<T>>;
// Define a structure to hold the weights in Eigen matrices
struct EigenModel {
aligned_vector<Eigen::MatrixXd> weight_matrices;
aligned_vector<Eigen::VectorXd> biases;
std::vector<Eigen::MatrixXd> weight_matrices;
std::vector<Eigen::VectorXd> biases;
};
struct TrainingData {
@ -46,8 +41,6 @@ struct TrainingData {
int Python_Keras_setup(std::string functions_file_path);
void set_ai_surrogate_runtime_params(RInsidePOET& R, RuntimeParameters& params, InitialList& init_list);
void Python_finalize(std::mutex* Eigen_model_mutex, std::mutex* training_data_buffer_mutex,
std::condition_variable* training_data_buffer_full, bool* start_training, bool* end_training);
@ -66,9 +59,9 @@ int Python_Keras_training_thread(EigenModel* Eigen_model,
bool* start_training, bool* end_training,
const RuntimeParameters& params);
void Python_Keras_set_weights_as_Eigen(EigenModel& eigen_model);
void update_weights(EigenModel* model, const std::vector<std::vector<std::vector<double>>>& weights);
Eigen::MatrixXd eigen_inference_batched(const Eigen::Ref<Eigen::MatrixXd>& input_batch, const EigenModel& model);
std::vector<std::vector<std::vector<double>>> Python_Keras_get_weights();
std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vector<double>> x, int batch_size,
std::mutex* Eigen_model_mutex);
@ -76,7 +69,6 @@ std::vector<double> Eigen_predict(const EigenModel& model, std::vector<std::vect
// Otherwise, define the necessary stubs
#else
inline void Python_Keras_setup(std::string functions_file_path){}
inline void set_ai_surrogate_runtime_params(RInsidePOET&, RuntimeParameters&, InitialList&){}
inline void Python_finalize(std::mutex*, std::mutex*, std::condition_variable*, bool*, bool*){}
inline void Python_Keras_load_model(std::string model_file_path){}
inline void training_data_buffer_append(std::vector<std::vector<double>>&, std::vector<std::vector<double>>&){}
@ -85,7 +77,9 @@ inline int Python_Keras_training_thread(EigenModel*, std::mutex*,
TrainingData*, std::mutex*,
std::condition_variable*, RuntimeParameters&,
bool*, bool*){return {};}
inline void Python_Keras_set_weights_as_Eigen(EigenModel&){}
inline void update_weights(EigenModel*, const std::vector<std::vector<std::vector<double>>>&){return {};}
inline std::vector<std::vector<std::vector<double>>> Python_Keras_get_weights(){return {};}
inline std::vector<double> Eigen_predict(const EigenModel&, std::vector<std::vector<double>>, int, std::mutex*){return {};}
#endif
} // namespace poet

53
src/poet.cpp
View File

@ -305,8 +305,23 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters &params,
&training_data_buffer_full, &start_training, &end_training,
params);
if (!params.use_Keras_predictions) {
// Initialize Eigen model for custom inference function
MSG("AI: Use custom C++ prediction function");
Python_Keras_set_weights_as_Eigen(Eigen_model);
// Get Keras weights from Python
std::vector<std::vector<std::vector<double>>> cpp_weights = Python_Keras_get_weights();
// Set model size
size_t num_layers = cpp_weights.size() / 2;
Eigen_model.weight_matrices.resize(num_layers);
Eigen_model.biases.resize(num_layers);
for (size_t i = 0; i < cpp_weights.size(); i += 2) {
size_t rows = cpp_weights[i][0].size();
size_t cols = cpp_weights[i].size();
Eigen_model.weight_matrices[i / 2].resize(rows, cols);
size_t bias_size = cpp_weights[i + 1][0].size();
Eigen_model.biases[i / 2].resize(bias_size);
}
// Set initial model weights
update_weights(&Eigen_model, cpp_weights);
}
MSG("AI: Surrogate model initialized");
}
@ -358,6 +373,9 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters &params,
R.parseEval(std::string("predictions_scaled <- ") +
"set_field(TMP, ai_surrogate_species, field_nrow, ai_surrogate_species, byrow = TRUE)");
R.parseEval("predictions <- postprocess(predictions_scaled)");
R.parseEval("print(head(predictions))");
// Validate prediction and write valid predictions to chem field
MSG("AI: Validate");
R.parseEval("validity_vector <- validate_predictions(predictors, predictions)");
@ -626,7 +644,38 @@ int main(int argc, char *argv[]) {
throw std::runtime_error("AI surrogate input script must contain a value for model_file_path");
}
set_ai_surrogate_runtime_params(R, run_params, init_list);
/* AI surrogate training and inference parameters. (Can be set by declaring a
variable of the same name in one of the the R input scripts)*/
run_params.Keras_training_always_use_CPU = false; // Default will use GPU if detected
run_params.Keras_training_always_use_CPU = false; // Default will use GPU if detected
run_params.use_Keras_predictions = false; // Default inference function is custom C++ / Eigen implementation
run_params.batch_size = 2560; // default value determined in test on the UP Turing cluster
run_params.training_epochs = 20; //
run_params.training_data_size = init_list.getDiffusionInit().n_rows *
init_list.getDiffusionInit().n_cols; // Default value is number of cells in field
run_params.save_model_path = ""; // Model is only saved if a path is set in the input field
if (Rcpp::as<bool>(R.parseEval("exists(\"batch_size\")"))) {
run_params.batch_size = R["batch_size"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"training_epochs\")"))) {
run_params.training_epochs = R["training_epochs"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"training_data_size\")"))) {
run_params.training_data_size = R["training_data_size"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"use_Keras_predictions\")"))) {
run_params.use_Keras_predictions = R["use_Keras_predictions"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"Keras_predictions_always_use_CPU\")"))) {
run_params.Keras_predictions_always_use_CPU = R["Keras_predictions_always_use_CPU"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"Keras_training_always_use_CPU\")"))) {
run_params.Keras_training_always_use_CPU = R["Keras_training_always_use_CPU"];
}
if (Rcpp::as<bool>(R.parseEval("exists(\"save_model_path\")"))) {
run_params.save_model_path = Rcpp::as<std::string>(R["save_model_path"]);
std::cout << "AI: Model will be saved as \"" << run_params.save_model_path << "\"" << std::endl;
}
MSG("AI: Initialize Python for AI surrogate functions");
std::string python_keras_file = std::string(SRC_DIR) +