From 105e091f43dd682922cbb1c0b38dedff683505ab Mon Sep 17 00:00:00 2001 From: moneta Date: Tue, 27 Jan 2026 16:14:13 +0100 Subject: [PATCH] [tmva][sofie] Add support for parsing external ONNX data file for weights Support external_data when parsing the initialized tensor from an ONNX file To test the feture, run tutorial RMVA_SOFIE_ONNX.py changing when exporting to the ONNX file from external_data=False to external_data=True --- tmva/sofie/inc/TMVA/SOFIE_common.hxx | 12 +- .../inc/TMVA/RModelParser_ONNX.hxx | 17 ++ tmva/sofie_parsers/src/RModelParser_ONNX.cxx | 211 +++++++++++------- 3 files changed, 147 insertions(+), 93 deletions(-) diff --git a/tmva/sofie/inc/TMVA/SOFIE_common.hxx b/tmva/sofie/inc/TMVA/SOFIE_common.hxx index 4d0da2304a729..7bac4af5c0307 100644 --- a/tmva/sofie/inc/TMVA/SOFIE_common.hxx +++ b/tmva/sofie/inc/TMVA/SOFIE_common.hxx @@ -283,16 +283,8 @@ public: for (std::size_t item : fShape) { fSize *= static_cast(item); } - switch (fType) { - case ETensorType::FLOAT: fSize *= sizeof(float); break; - case ETensorType::DOUBLE: fSize *= sizeof(double); break; - case ETensorType::INT32: fSize *= sizeof(int32_t); break; - case ETensorType::INT64: fSize *= sizeof(int64_t); break; - case ETensorType::BOOL: fSize *= sizeof(bool); break; - default: - throw std::runtime_error("TMVA::SOFIE doesn't yet supports serialising data-type " + - ConvertTypeToString(fType)); - } + // get size in bytes + fSize *= GetTypeSize(fType); fPersistentData = static_cast(fData.get()); } void CastPersistentToShared() diff --git a/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx b/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx index ed42f7381693b..395a1c7db5eeb 100644 --- a/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx +++ b/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx @@ -6,12 +6,14 @@ #include #include #include +#include // forward declaration namespace onnx { class NodeProto; class GraphProto; class ModelProto; +class TensorProto; } // namespace onnx namespace TMVA { @@ -42,6 +44,11 @@ private: // List of fused operators storing as key the second operator and a value a pair of fusion type and parent operator std::map> fFusedOperators; + // weight data file + std::ifstream fDataFile; + // filename of model + std::string fDataFileName; + public: // Register an ONNX operator @@ -80,6 +87,10 @@ public: std::unique_ptr LoadModel(const std::string &filename); std::unique_ptr LoadModel(std::istream &input); + std::shared_ptr GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_length, ETensorType type ); + +public: + RModelParser_ONNX() noexcept; RModel Parse(std::string const &filename, bool verbose = false); @@ -88,6 +99,12 @@ public: // check the model for missing operators - return false in case some operator implementation is missing bool CheckModel(std::string filename, bool verbose = false); + //set external data full path (needed if external data are not stored in the default modelName.onnx.data) + // call this function before parsing + void SetExternalDataFile(const std::string & dataFileName) { + fDataFileName = dataFileName; + } + ~RModelParser_ONNX(); }; diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx index 25800fb4d597c..aa444c05d5f76 100644 --- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx +++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx @@ -120,62 +120,130 @@ struct ExtractDataFromTP { // trait function to extract data from TensorProto template<> struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * tensor, void * data) { + static void Copy(onnx::TensorProto * tensor, void * data, int length) { + if (tensor->float_data_size() != length) + throw std::runtime_error("TMVA::SOFIE - Failed to read float initialized tensor - actual size is " + std::to_string(tensor->float_data_size())); tensor->mutable_float_data()->ExtractSubrange(0, tensor->float_data_size(), static_cast(data)); } }; template<> struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * tensor, void * data) { + static void Copy(onnx::TensorProto * tensor, void * data, int length) { + if (tensor->double_data_size() != length) + throw std::runtime_error("TMVA::SOFIE - Failed to read double initialized tensor - actual size is " + std::to_string(tensor->double_data_size())); tensor->mutable_double_data()->ExtractSubrange(0, tensor->double_data_size(), static_cast(data)); } }; template<> struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * tensor, void * data) { + static void Copy(onnx::TensorProto * tensor, void * data, int length) { + if (tensor->int32_data_size() != length) + throw std::runtime_error("TMVA::SOFIE - Failed to read int32 initialized tensor - actual size is " + std::to_string(tensor->int32_data_size())); tensor->mutable_int32_data()->ExtractSubrange(0, tensor->int32_data_size(), static_cast(data)); } }; template<> struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * tensor, void * data) { + static void Copy(onnx::TensorProto * tensor, void * data, int length) { + if (tensor->int64_data_size() != length) + throw std::runtime_error("TMVA::SOFIE - Failed to read int64 initialized tensor - actual size is " + std::to_string(tensor->int64_data_size())); tensor->mutable_int64_data()->ExtractSubrange(0, tensor->int64_data_size(), static_cast(data)); } }; -template<> -struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * , void * ) { - throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in UINT8 not supported"); - } -}; -template<> -struct ExtractDataFromTP { - static void Copy(onnx::TensorProto * , void *) { - throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in BOOL not supported"); - } -}; -template -std::shared_ptr GetInitializedTensorData(onnx::TensorProto * tensorproto, size_t length) { - std::shared_ptr data(malloc(length * sizeof(T)), free); - if (!tensorproto->raw_data().empty()) { +std::shared_ptr RModelParser_ONNX::GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_size, ETensorType tensor_type) +{ + + std::shared_ptr data(malloc(tensor_size), free); + + // check if initialized tensors are stored internally + if (tensorproto->data_location() != onnx::TensorProto::EXTERNAL) { + if (tensorproto->raw_data().size() > 0) { + if (tensorproto->raw_data().size() != tensor_size) + throw std::runtime_error("TMVA::SOFIE - Failed to read raw data of initialized tensor - actual raw size is " + + std::to_string(tensorproto->raw_data().size())); + #ifdef R__BYTESWAP - std::memcpy(data.get(), tensorproto->raw_data().c_str(), length * sizeof(T)); + // R__BYTESWAP is defined for little-endian architectures (most common ones) + std::memcpy(data.get(), tensorproto->raw_data().c_str(), tensor_size); #else - for (std::size_t k = 0; k < length; ++k) - (reinterpret_cast::value_type *>(data.get()))[k] = - RByteSwap::bswap((reinterpret_cast::value_type *>(tensorproto->raw_data().c_str()))[k]); + // big-endian architectures - need to swap bytes + for (std::size_t k = 0; k < tensor_size; ++k) + (reinterpret_cast::value_type *>(data.get()))[k] = + RByteSwap::bswap((reinterpret_cast::value_type *>( + tensorproto->raw_data().c_str()))[k]); #endif - } else { - ExtractDataFromTP::Copy(tensorproto, data.get()); + } else { + // case tensor data are stored as specific types and now in raw_data + switch (tensor_type) { + case ETensorType::FLOAT: { + ExtractDataFromTP::Copy(tensorproto, data.get(), tensor_size/ 4); + break; + } + case ETensorType::DOUBLE: { + ExtractDataFromTP::Copy(tensorproto, data.get(), tensor_size/ 8); + break; + } + case ETensorType::INT32: { + ExtractDataFromTP::Copy(tensorproto, data.get(), tensor_size/ 4); + break; + } + case ETensorType::INT64: { + ExtractDataFromTP::Copy(tensorproto, data.get(), tensor_size/ 8); + break; + } + case ETensorType::BOOL: { + throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in BOOL not supported"); + break; + } + case ETensorType::UINT8: { + throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in UINT8 not supported"); + break; + } + default: + throw std::runtime_error("Data type " + ConvertTypeToString(tensor_type) + " in weight tensor is not supported!\n"); + } + } + + } else { + // case of external data + if (fVerbose) + std::cout << "Initialized data are stored externally in file " << fDataFileName; + + // read now tensor from file + std::string location; + size_t offset = 0, buffer_size = 0; + + for (const auto &kv : tensorproto->external_data()) { + if (kv.key() == "location") location = kv.value(); + else if (kv.key() == "offset") offset = std::stoull(kv.value()); + else if (kv.key() == "length") buffer_size = std::stoull(kv.value()); + } + if (fVerbose) + std::cout << " at location " << location << " offset " << offset << " and with length " << buffer_size << std::endl; + + if (buffer_size != tensor_size) + throw std::runtime_error("TMVA::SOFIE ONNX : invalid stored data size vs tensor size"); + + // open the data file if needed + if (!fDataFile.is_open()) { + fDataFile.open(fDataFileName, std::ios::binary); + if (!fDataFile.is_open()) + throw std::runtime_error("TMVA::SOFIE ONNX: error reading external weight ONNX data file " + fDataFileName); + } + + fDataFile.seekg(offset); + fDataFile.read(reinterpret_cast(data.get()), buffer_size); } + return data; } + // Constructor of the parser RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_unique()) { // Register operators @@ -421,6 +489,8 @@ RModel RModelParser_ONNX::Parse(std::string const &filename, bool verbose) filename_nodir = (filename.substr(isep + 1, filename.length() - isep)); } + if (fDataFileName.empty() ) fDataFileName = filename + ".data"; + RModel rmodel(filename_nodir, parsetime); ParseONNXGraph(rmodel, graph, filename_nodir); return rmodel; @@ -621,74 +691,49 @@ void RModelParser_ONNX::ParseONNXGraph(RModel & rmodel, const onnx::GraphProto & for (int i = 0; i < graph.initializer_size(); i++) { onnx::TensorProto *tensorproto = const_cast(&graph.initializer(i)); std::vector shape; - std::size_t fLength = 1; + std::size_t tensor_length = 1; for (int j = 0; j < tensorproto->dims_size(); j++) { shape.push_back(tensorproto->dims(j)); - fLength *= tensorproto->dims(j); + tensor_length *= tensorproto->dims(j); } // in case of scalars keep an empty shape but with length =1 - std::string input_name = graph.initializer(i).name(); + std::string tensor_name = graph.initializer(i).name(); if (verbose) - std::cout << "\t initializer " << i << " name " << input_name << " type " << graph.initializer(i).data_type() - << std::endl; + std::cout << "\t initializer " << i << " name " << tensor_name << " type " << graph.initializer(i).data_type() + << " and length " << tensor_length << std::endl; + // register also the initialized tensors auto tensor_type = static_cast(graph.initializer(i).data_type()); - RegisterTensorType(input_name, tensor_type); - - switch (tensor_type) { - case ETensorType::FLOAT: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add FLOAT initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::FLOAT, shape, data); - allInitializedTensors[input_name] = i; - break; - } - case ETensorType::DOUBLE: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add DOUBLE initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::DOUBLE, shape, data); - allInitializedTensors[input_name] = i; - break; - } - case ETensorType::INT32: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add INT32 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::INT32, shape, data); - allInitializedTensors[input_name] = i; - break; - } - case ETensorType::INT64: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add INT64 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::INT64, shape, data); - allInitializedTensors[input_name] = i; - break; - } - case ETensorType::BOOL: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add BOOL initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::BOOL, shape, data); - allInitializedTensors[input_name] = i; - if (verbose) - std::cout << "Bool initialized data: " << ConvertValuesToString(fLength,reinterpret_cast(data.get(),10)) << std::endl; - break; - } - case ETensorType::UINT8: { - std::shared_ptr data = GetInitializedTensorData(tensorproto, fLength); - if (verbose) std::cout << "add UINT8 initialized tensor " << input_name << " shape " << ConvertShapeToString(shape) << std::endl; - rmodel.AddInitializedTensor(input_name, ETensorType::UINT8, shape, data); - allInitializedTensors[input_name] = i; - if (verbose) - std::cout << "uint8 initialized data: " << ConvertValuesToString(fLength,reinterpret_cast(data.get(),10)) << std::endl; - break; - } - default: - throw std::runtime_error("Data type in weight tensor " + graph.initializer(i).name() + " not supported!\n"); + RegisterTensorType(tensor_name, tensor_type); + + std::shared_ptr data = GetInitializedTensorData(tensorproto, tensor_length * GetTypeSize(tensor_type), tensor_type); + rmodel.AddInitializedTensor(tensor_name, tensor_type, shape, data); + allInitializedTensors[tensor_name] = i; + + if (verbose) { + std::cout << "add initialized tensor " << tensor_name << "with shape " << ConvertShapeToString(shape) << "and "; + if (tensor_type == ETensorType::FLOAT) { + std::cout << " float data: "; + for (int j = 0; j < std::min(int(tensor_length),3); j++) std::cout << static_cast(data.get())[j] << " "; + } + else if (tensor_type == ETensorType::INT64) { + std::cout << " int64 data: "; + for (int j = 0; j < std::min(int(tensor_length),3); j++) std::cout << static_cast(data.get())[j] << " "; + } + else if (tensor_type == ETensorType::UINT8) { + std::cout << " uint8 data: "; + for (int j = 0; j < std::min(int(tensor_length),3); j++) std::cout << static_cast(data.get())[j] << " "; + } + else if (tensor_type == ETensorType::BOOL) { + std::cout << " Boolean data: "; + for (int j = 0; j < std::min(int(tensor_length),3); j++) std::cout << static_cast(data.get())[j] << " "; + } + std::cout << std::endl; } - } + } // end initializer list // Initial operator order if (verbose) {