From 8691370abb7f46bbcc0970a835748b8e853cf0e5 Mon Sep 17 00:00:00 2001 From: straile Date: Tue, 15 Oct 2024 11:17:30 +0200 Subject: [PATCH] Eigen model works --- .../AI_Python_functions/keras_AI_surrogate.py | 3 + .../SurrogateModels/AI_functions.cpp | 292 +++++++++--------- .../SurrogateModels/AI_functions.hpp | 22 +- src/poet.cpp | 53 +++- 4 files changed, 205 insertions(+), 165 deletions(-) diff --git a/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py b/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py index 879029bae..a61c255a3 100644 --- a/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py +++ b/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py @@ -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() diff --git a/src/Chemistry/SurrogateModels/AI_functions.cpp b/src/Chemistry/SurrogateModels/AI_functions.cpp index 5f25811f3..ef94408a1 100644 --- a/src/Chemistry/SurrogateModels/AI_functions.cpp +++ b/src/Chemistry/SurrogateModels/AI_functions.cpp @@ -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(R.parseEval("exists(\"batch_size\")"))) { - params.batch_size = R["batch_size"]; - } - if (Rcpp::as(R.parseEval("exists(\"training_epochs\")"))) { - params.training_epochs = R["training_epochs"]; - } - if (Rcpp::as(R.parseEval("exists(\"training_data_size\")"))) { - params.training_data_size = R["training_data_size"]; - } - if (Rcpp::as(R.parseEval("exists(\"use_Keras_predictions\")"))) { - params.use_Keras_predictions = R["use_Keras_predictions"]; - } - if (Rcpp::as(R.parseEval("exists(\"Keras_predictions_always_use_CPU\")"))) { - params.Keras_predictions_always_use_CPU = R["Keras_predictions_always_use_CPU"]; - } - if (Rcpp::as(R.parseEval("exists(\"Keras_training_always_use_CPU\")"))) { - params.Keras_training_always_use_CPU = R["Keras_training_always_use_CPU"]; - } - if (Rcpp::as(R.parseEval("exists(\"save_model_path\")"))) { - params.save_model_path = Rcpp::as(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 Python_Keras_predict(std::vector> 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& 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,46 +182,51 @@ std::vector Python_Keras_predict(std::vector> 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 Eigen_predict(const EigenModel& model, std::vector> x, int batch_size, +std::vector Eigen_predict(const EigenModel& model, const std::vector>& 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]; + // 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]; + } } - } - std::vector result; - result.reserve(num_samples * num_features); - - 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())); - } - int num_batches = std::ceil(static_cast(num_samples) / batch_size); + std::vector result; + result.reserve(num_samples); - Eigen_model_mutex->lock(); - 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); + 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())); + } - result.insert(result.end(), batch_data.data(), batch_data.data() + batch_data.size()); - } - Eigen_model_mutex->unlock(); - return result; + int num_batches = std::ceil(static_cast(num_samples) / batch_size); + + std::lock_guard 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 = 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()); + } + + return result; } + /** * @brief Appends data from one matrix (column major std::vector>) 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& 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>>& 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]; + } + } + // 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]; + } } - // 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 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>> 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>> 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"); + 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."); + } - PyArrayObject* weight_np = reinterpret_cast(weight_array); - if (PyArray_NDIM(weight_np) != 2) throw std::runtime_error("Weight array is not 2-dimensional"); + PyArrayObject* weight_np = reinterpret_cast(py_weight_array); + int dtype = PyArray_TYPE(weight_np); - // 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"); - } + 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); + + // Create a container for the matrix + std::vector> weight_matrix(num_rows, std::vector(num_cols)); - // Cast the data correctly depending on the type - void* weight_data = PyArray_DATA(weight_np); - int num_rows = PyArray_DIM(weight_np, 0); - int num_cols = PyArray_DIM(weight_np, 1); + // Handle different data types + if (dtype == NPY_FLOAT32) { + float* weight_data_float = static_cast(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(weight_data_float[r * num_cols + c]); + } + } + } else if (dtype == NPY_DOUBLE) { + double* weight_data_double = static_cast(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."); + } - Eigen::MatrixXd weight_matrix; - if (dtype == NPY_FLOAT32) { - // Handle float32 array from NumPy - float* weight_data_float = static_cast(weight_data); - weight_matrix = Eigen::Map>( - weight_data_float, num_rows, num_cols).cast().transpose(); - } else if (dtype == NPY_DOUBLE) { - // Handle double (float64) array from NumPy - double* weight_data_double = static_cast(weight_data); - weight_matrix = Eigen::Map>( - weight_data_double, num_rows, num_cols).transpose(); - } + cpp_weights.push_back(weight_matrix); - // Get bias vector - PyObject* bias_array = PyList_GetItem(py_weights_list, i + 1); - PyArrayObject* bias_np = reinterpret_cast(bias_array); - if (PyArray_NDIM(bias_np) != 1) throw std::runtime_error("Bias array is not 1-dimensional"); + } else if (PyArray_NDIM(weight_np) == 1) { + // Handle 1D bias vectors + int num_elements = PyArray_DIM(weight_np, 0); + + // Create a container for the vector + std::vector> bias_vector(1, std::vector(num_elements)); - // 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; + // Handle different data types + if (dtype == NPY_FLOAT32) { + float* bias_data_float = static_cast(PyArray_DATA(weight_np)); + for (int j = 0; j < num_elements; ++j) { + bias_vector[0][j] = static_cast(bias_data_float[j]); + } + } else if (dtype == NPY_DOUBLE) { + double* bias_data_double = static_cast(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."); + } - if (bias_dtype == NPY_FLOAT32) { - float* bias_data_float = static_cast(bias_data); - bias_vector = Eigen::Map>(bias_data_float, bias_size).cast(); - } else if (bias_dtype == NPY_DOUBLE) { - double* bias_data_double = static_cast(bias_data); - bias_vector = Eigen::Map(bias_data_double, bias_size); - } - - // 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; } diff --git a/src/Chemistry/SurrogateModels/AI_functions.hpp b/src/Chemistry/SurrogateModels/AI_functions.hpp index 69a9f4c20..ad8457559 100644 --- a/src/Chemistry/SurrogateModels/AI_functions.hpp +++ b/src/Chemistry/SurrogateModels/AI_functions.hpp @@ -17,8 +17,6 @@ #include #include #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 -using aligned_vector = std::vector>; -// Define a structure to hold the weights in Eigen matrices + struct EigenModel { - aligned_vector weight_matrices; - aligned_vector biases; + std::vector weight_matrices; + std::vector 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>>& weights); -Eigen::MatrixXd eigen_inference_batched(const Eigen::Ref& input_batch, const EigenModel& model); +std::vector>> Python_Keras_get_weights(); std::vector Eigen_predict(const EigenModel& model, std::vector> x, int batch_size, std::mutex* Eigen_model_mutex); @@ -76,7 +69,6 @@ std::vector Eigen_predict(const EigenModel& model, std::vector>&, std::vector>&){} @@ -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>>&){return {};} +inline std::vector>> Python_Keras_get_weights(){return {};} inline std::vector Eigen_predict(const EigenModel&, std::vector>, int, std::mutex*){return {};} #endif } // namespace poet diff --git a/src/poet.cpp b/src/poet.cpp index 864b2f17b..6779885ea 100644 --- a/src/poet.cpp +++ b/src/poet.cpp @@ -305,8 +305,23 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters ¶ms, &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>> 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 ¶ms, 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(R.parseEval("exists(\"batch_size\")"))) { + run_params.batch_size = R["batch_size"]; + } + if (Rcpp::as(R.parseEval("exists(\"training_epochs\")"))) { + run_params.training_epochs = R["training_epochs"]; + } + if (Rcpp::as(R.parseEval("exists(\"training_data_size\")"))) { + run_params.training_data_size = R["training_data_size"]; + } + if (Rcpp::as(R.parseEval("exists(\"use_Keras_predictions\")"))) { + run_params.use_Keras_predictions = R["use_Keras_predictions"]; + } + if (Rcpp::as(R.parseEval("exists(\"Keras_predictions_always_use_CPU\")"))) { + run_params.Keras_predictions_always_use_CPU = R["Keras_predictions_always_use_CPU"]; + } + if (Rcpp::as(R.parseEval("exists(\"Keras_training_always_use_CPU\")"))) { + run_params.Keras_training_always_use_CPU = R["Keras_training_always_use_CPU"]; + } + if (Rcpp::as(R.parseEval("exists(\"save_model_path\")"))) { + run_params.save_model_path = Rcpp::as(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) +