From a0fe891f99f933dc20661acf3a4b34f73914b4a7 Mon Sep 17 00:00:00 2001 From: straile Date: Mon, 7 Oct 2024 22:01:29 +0200 Subject: [PATCH] Eigen inference but slow --- .../AI_Python_functions/keras_AI_surrogate.py | 10 +- .../SurrogateModels/AI_functions.cpp | 249 ++++++++++++++++-- .../SurrogateModels/AI_functions.hpp | 26 +- src/poet.cpp | 23 +- src/poet.hpp.in | 1 + 5 files changed, 280 insertions(+), 29 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 be6b3f079..ebc772ec2 100644 --- a/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py +++ b/src/Chemistry/SurrogateModels/AI_Python_functions/keras_AI_surrogate.py @@ -17,4 +17,12 @@ 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) \ No newline at end of file + print("Prediction from Python", flush=True) + print(prediction, flush=True) + return np.array(prediction, dtype=np.float64) + +def get_weights(model): + weights = model.get_weights() + # Convert all weight and bias arrays to float64 (double precision) + weights_as_double = [w.astype(np.float64) for w in weights] + return weights_as_double \ No newline at end of file diff --git a/src/Chemistry/SurrogateModels/AI_functions.cpp b/src/Chemistry/SurrogateModels/AI_functions.cpp index b874e6b9e..35cb43344 100644 --- a/src/Chemistry/SurrogateModels/AI_functions.cpp +++ b/src/Chemistry/SurrogateModels/AI_functions.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "AI_functions.hpp" @@ -16,6 +17,7 @@ namespace poet { * @brief Loads the Python interpreter and functions * @param functions_file_path Path to the Python file where the AI surrogate * functions are defined + * @return 0 if function was succesful */ int Python_Keras_setup(std::string functions_file_path) { // Initialize Python functions @@ -34,6 +36,7 @@ int Python_Keras_setup(std::string functions_file_path) { * @brief Loads the user-supplied Keras model * @param model_file_path Path to a .keras file that the user must supply as * a variable "model_file_path" in the R input script + * @return 0 if function was succesful */ int Python_Keras_load_model(std::string model_file_path) { // Initialize Keras model @@ -50,6 +53,7 @@ int Python_Keras_load_model(std::string model_file_path) { * @brief Converts the std::vector 2D matrix representation of a POET Field object to a numpy array * for use in the Python AI surrogate functions * @param field 2D-Matrix with the content of a Field object + * @return Numpy representation of the input vector */ PyObject* vector_to_numpy_array(const std::vector>& field) { // import numpy and numpy array API @@ -71,46 +75,44 @@ PyObject* vector_to_numpy_array(const std::vector>& field) { } /** - * @brief Converts a Pyton matrix object to a std::vector vector 2D matrix + * @brief Converts a Pyton matrix object to a std::vector vector * @param py_matrix Pyobject that must be a 2D matrix + * @result Vector that can be used similar to the return value of the Field object + * Field.AsVector() method. */ -std::vector> numpy_array_to_vector(PyObject* pyArray) { - std::vector> result; - if (!PyArray_Check(pyArray)) { +std::vector numpy_array_to_vector(PyObject* py_array) { + std::vector result; + if (!PyArray_Check(py_array)) { std::cerr << "The model's predictions are not a numpy array." << std::endl; return result; } // Cast generic PyObject to PyArrayObject - PyArrayObject* npArray = reinterpret_cast(pyArray); + PyArrayObject* np_array = reinterpret_cast(py_array); // Get shape - int numDims = PyArray_NDIM(npArray); - npy_intp* shape = PyArray_SHAPE(npArray); - + int numDims = PyArray_NDIM(np_array); + npy_intp* shape = PyArray_SHAPE(np_array); if (numDims != 2) { std::cerr << "The model's predictions are not a 2D matrix." << std::endl; return result; } // Copy data into std::vector format - double* data = static_cast(PyArray_DATA(npArray)); - result.resize(shape[0]); - for (npy_intp i = 0; i < shape[0]; ++i) { - result[i].resize(shape[1]); - for (npy_intp j = 0; j < shape[1]; ++j) { - result[i][j] = data[i * shape[1] + j]; - } - } + double* data = static_cast(PyArray_DATA(np_array)); + npy_intp size = PyArray_SIZE(np_array); + result.resize(size); + std::copy(data, data + size, result.begin()); return result; } - /** - * @brief Converts the 2D-Matrix represantation of a POET Field object to a numpy array - * for use in the Python AI surrogate functions - * @param field 2D-Matrix with the content of a Field object + * @brief Uses the Python Keras functions to calculate predictions from a neural network. + * @param x 2D-Matrix with the content of a Field object + * @param batch_size size for mini-batches that are used in the Keras model.predict() method + * @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> Python_keras_predict(std::vector> x, int batch_size) { +std::vector Python_keras_predict(std::vector> x, int batch_size) { // Prepare data for python PyObject* py_df_x = vector_to_numpy_array(x); @@ -127,7 +129,7 @@ std::vector> Python_keras_predict(std::vector> predictions = numpy_array_to_vector(py_predictions); + std::vector predictions = numpy_array_to_vector(py_predictions); // Clean up PyErr_Print(); // Ensure that python errors make it to stdout Py_DECREF(py_df_x); @@ -138,6 +140,207 @@ std::vector> Python_keras_predict(std::vector& input_batch, const EigenModel& model) { + Eigen::MatrixXd current_layer = input_batch; + std::cout << "input_batch ROWS " << current_layer.rows() << std::endl; + std::cout << "input_batch COLS " << current_layer.cols() << std::endl; + // Process all hidden layers + for (size_t layer = 0; layer < model.weight_matrices.size() - 1; ++layer) { + + std::cout << "WEIGHTS LAYER " << layer << std::endl; + std::cout << "WEIGHTS ROWS " << model.weight_matrices[layer].rows() << std::endl; + std::cout << "WEIGHTS COLS " << model.weight_matrices[layer].cols() << std::endl; + std::cout << "BIASES SIZE " << model.biases[layer].size() << std::endl; + + current_layer = (model.weight_matrices[layer] * current_layer); + std::cout << "MULTIPLIED" << std::endl; + 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 Converts the weights and biases from a Python Keras model to Eigen matrices + * @return The struct containing the model weights and biases as aligned Eigen matrices + */ +#include + +EigenModel Python_Keras_get_weights_as_Eigen() { + EigenModel eigen_model; + + PyObject* py_main_module = PyImport_AddModule("__main__"); + if (!py_main_module) throw std::runtime_error("Failed to import Python main module"); + + PyObject* py_global_dict = PyModule_GetDict(py_main_module); + if (!py_global_dict) throw std::runtime_error("Failed to get global dictionary"); + + PyObject* py_keras_model = PyDict_GetItemString(py_global_dict, "model"); + if (!py_keras_model) throw std::runtime_error("Failed to get Keras model from Python dictionary"); + + PyObject* py_get_weights_function = PyDict_GetItemString(py_global_dict, "get_weights"); + if (!py_get_weights_function) throw std::runtime_error("Failed to get Keras get_weights function"); + + PyObject* args = Py_BuildValue("(O)", py_keras_model); + if (!args) throw std::runtime_error("Failed to create argument tuple"); + + // Call Python function + PyObject* py_weights_list = PyObject_CallObject(py_get_weights_function, args); + Py_DECREF(args); // Cleanup args + + if (!py_weights_list) { + PyErr_Print(); // Print Python error + throw std::runtime_error("Failed to get weights from Keras model"); + } + + 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(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) { + std::cout << "DTYPE: SINGLE" << std::endl; + } + if (dtype == NPY_DOUBLE) { + std::cout << "DTYPE: DOUBLE" << std::endl; + } + if (dtype != NPY_FLOAT32 && dtype != NPY_DOUBLE) { + throw std::runtime_error("Unsupported data type for weight array"); + } + + // Cast the data correctly depending on the type + void* weight_data = PyArray_DATA(weight_np); + if (!weight_data) throw std::runtime_error("Failed to get weight data pointer"); + + int num_rows = PyArray_DIM(weight_np, 0); // NumPy's row count + int num_cols = PyArray_DIM(weight_np, 1); // NumPy's column count + + Eigen::MatrixXd weight_matrix; // Use double-precision 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(); + } + + // Get bias vector + PyObject* bias_array = PyList_GetItem(py_weights_list, i + 1); + if (!bias_array) throw std::runtime_error("Failed to get bias array"); + if (!PyArray_Check(bias_array)) throw std::runtime_error("Bias array is not a NumPy array"); + + PyArrayObject* bias_np = reinterpret_cast(bias_array); + if (PyArray_NDIM(bias_np) != 1) throw std::runtime_error("Bias array is not 1-dimensional"); + + // Check bias data type and cast accordingly + int bias_dtype = PyArray_TYPE(bias_np); + void* bias_data = PyArray_DATA(bias_np); + if (!bias_data) throw std::runtime_error("Failed to get bias data pointer"); + + int bias_size = PyArray_DIM(bias_np, 0); + Eigen::VectorXd bias_vector; // Use double-precision vector + + 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); + } + + // Clean up + Py_DECREF(py_weights_list); // Clean up Python list object + Py_DECREF(py_main_module); // Clean up Python module object + + return eigen_model; +} + + + + + +/** + * @brief Gets the model weights and biases of the Keras model and uses Eigen for fast inference + * @param x 2D-Matrix with the content of a Field object + * @param batch_size size for mini-batches that are used in the Keras model.predict() method + * @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) { + // Convert input data to Eigen matrix + const int num_samples = x[0].size(); + const int num_features = x.size(); + + std::cout << "num_samples " << num_samples << std::endl; + std::cout << "num_features " << num_features << std::endl; + + 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]; + if (i < 6) {std::cout << full_input_matrix.coeff(j, i) << " ";} + } + if (i < 6) {std::cout << std::endl;} + } + + + std::cout << "Eigen rows " << full_input_matrix.rows() << std::endl; + std::cout << "Eigen cols " << full_input_matrix.cols() << std::endl; + + std::vector result; + result.reserve(num_samples * model.biases.back().size()); + + int num_batches = std::ceil(static_cast(num_samples) / batch_size); + 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; + std::cout << "CREATE INPUT BATCH" << std::endl; + // Extract the current batch + Eigen::MatrixXd input_batch(num_features, current_batch_size); + input_batch = full_input_matrix.block(0, start_idx, num_features, current_batch_size); + // Perform inference on the batch + Eigen::MatrixXd output_batch(num_features, current_batch_size); + output_batch = eigen_inference_batched(input_batch, model); + for (int i = 0; i < num_samples; ++i) { + for (int j = 0; j < num_features; ++j) { + if (i < 6) {std::cout << output_batch.coeff(j, i) << " ";} + } + if (i < 6) {std::cout << std::endl;} + } + + //result.insert(result.end(), output_batch.data(), output_batch.data() + output_batch.size()); + } + + return result; +} } \ No newline at end of file diff --git a/src/Chemistry/SurrogateModels/AI_functions.hpp b/src/Chemistry/SurrogateModels/AI_functions.hpp index 60840d811..176ef9224 100644 --- a/src/Chemistry/SurrogateModels/AI_functions.hpp +++ b/src/Chemistry/SurrogateModels/AI_functions.hpp @@ -17,6 +17,12 @@ #include #include +// PhreeqC definition of pi clashes with Eigen macros so we have to temporarily undef it +#pragma push_macro("pi") +#undef pi +#include +#pragma pop_macro("pi") + namespace poet { int Python_Keras_setup(std::string functions_file_path); @@ -25,11 +31,27 @@ int Python_Keras_load_model(std::string model_file_path); PyObject* vector_to_numpy_array(const std::vector>& field); -std::vector> numpy_array_to_vector(PyObject* py_matrix); +std::vector numpy_array_to_vector(PyObject* py_matrix); -std::vector> Python_keras_predict(std::vector> x, int batch_size); +std::vector Python_keras_predict(std::vector> x, int batch_size); + +// 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; +}; + +EigenModel Python_Keras_get_weights_as_Eigen(); + +Eigen::MatrixXd Eigen_batched_inference(const Eigen::Ref& input_batch, const EigenModel& model); + +std::vector Eigen_predict(const EigenModel& model, std::vector> x, int batch_size); //int Python_keras_train(Field &x, Field &y, Field &x_val, Field &y_val, int batch_size, std::string pid) + } // namespace poet #endif // AI_FUNCTIONS_HPP \ No newline at end of file diff --git a/src/poet.cpp b/src/poet.cpp index 94134914a..a1e34adfe 100644 --- a/src/poet.cpp +++ b/src/poet.cpp @@ -275,6 +275,10 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters ¶ms, /* Iteration Count is dynamic, retrieving value from R (is only needed by * master for the following loop) */ uint32_t maxiter = params.timesteps.size(); + + /* Stores the weights and biases of the AI surrogate model + * for use in an inference function with Eigen */ + EigenModel Eigen_model; if (params.print_progress) { chem.setProgressBarPrintout(true); @@ -314,12 +318,25 @@ static Rcpp::List RunMasterLoop(RInsidePOET &R, const RuntimeParameters ¶ms, MSG("AI Preprocessing"); R.parseEval("predictors_scaled <- preprocess(predictors)"); - // Predict MSG("AI: Predict"); - R["predictions_scaled"] = Rcpp::wrap(Python_keras_predict(R["predictors_scaled"], params.batch_size)); - R.parseEval("setNames(predictions_scaled, TMP_POPS)"); + R["TMP"] = Python_keras_predict(R["predictors_scaled"], params.batch_size); + R.parseEval("predictions_scaled <- matrix(TMP, nrow=nrow(predictors), byrow = TRUE)"); + R.parseEval("predictions_scaled <- setNames(data.frame(predictions_scaled), colnames(predictors))"); + MSG("Prediction from Python sent to R") + R.parseEval("print(head(predictions_scaled))"); + + EigenModel Eigen_model = Python_Keras_get_weights_as_Eigen(); + MSG("EIGEN WEIGHTS SUCESFULLY IMPORTED"); + R["TMP"] = Eigen_predict(Eigen_model, R["predictors_scaled"], params.batch_size); + R.parseEval("predictions_scaled <- matrix(TMP, nrow=nrow(predictors), byrow = TRUE)"); + MSG("EIGEN SUCESFULLY PREDICTED"); + R.parseEval("predictions_scaled <- setNames(data.frame(predictions_scaled), colnames(predictors))"); + + MSG("Prediction from Eigen sent to R") + R.parseEval("print(head(predictions_scaled))"); + // after this comes old R code! // Apply postprocessing MSG("AI Postprocesing"); diff --git a/src/poet.hpp.in b/src/poet.hpp.in index 261ed428c..a1d89fa14 100644 --- a/src/poet.hpp.in +++ b/src/poet.hpp.in @@ -67,4 +67,5 @@ struct RuntimeParameters { std::uint32_t interp_bucket_entries; bool use_ai_surrogate = false; + int batch_size; };