Skip to content

[tmva][sofie] Add support for parsing external ONNX data file for weights #21048

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lmoneta merged 1 commit into from
Apr 11, 2026
Merged

[tmva][sofie] Add support for parsing external ONNX data file for weights #21048

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 2 additions & 10 deletions tmva/sofie/inc/TMVA/SOFIE_common.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,16 +283,8 @@ public:
for (std::size_t item : fShape) {
fSize *= static_cast<int>(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<char *>(fData.get());
}
void CastPersistentToShared()
Expand Down
17 changes: 17 additions & 0 deletions tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,14 @@
#include <memory>
#include <functional>
#include <unordered_map>
#include <fstream>

// forward declaration
namespace onnx {
class NodeProto;
class GraphProto;
class ModelProto;
class TensorProto;
} // namespace onnx

namespace TMVA {
Expand Down Expand Up @@ -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<int, std::pair<EFusedOp, int>> fFusedOperators;

// weight data file
std::ifstream fDataFile;
// filename of model
std::string fDataFileName;


public:
// Register an ONNX operator
Expand Down Expand Up @@ -80,6 +87,10 @@ public:
std::unique_ptr<onnx::ModelProto> LoadModel(const std::string &filename);
std::unique_ptr<onnx::ModelProto> LoadModel(std::istream &input);

std::shared_ptr<void> GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_length, ETensorType type );

public:

RModelParser_ONNX() noexcept;

RModel Parse(std::string const &filename, bool verbose = false);
Expand All @@ -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();
};

Expand Down
211 changes: 128 additions & 83 deletions tmva/sofie_parsers/src/RModelParser_ONNX.cxx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,62 +120,130 @@ struct ExtractDataFromTP {
// trait function to extract data from TensorProto
template<>
struct ExtractDataFromTP<float> {
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<float *>(data));
}
};
template<>
struct ExtractDataFromTP<double> {
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<double *>(data));
}
};
template<>
struct ExtractDataFromTP<int32_t> {
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<int32_t *>(data));
}
};
template<>
struct ExtractDataFromTP<int64_t> {
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<int64_t *>(data));
}
};
template<>
struct ExtractDataFromTP<uint8_t> {
static void Copy(onnx::TensorProto * , void * ) {
throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in UINT8 not supported");
}
};
template<>
struct ExtractDataFromTP<bool> {
static void Copy(onnx::TensorProto * , void *) {
throw std::runtime_error("TMVA::SOFIE - ExtractData from TP in BOOL not supported");
}
};
template<typename T>
std::shared_ptr<void> GetInitializedTensorData(onnx::TensorProto * tensorproto, size_t length) {
std::shared_ptr<void> data(malloc(length * sizeof(T)), free);

if (!tensorproto->raw_data().empty()) {
std::shared_ptr<void> RModelParser_ONNX::GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_size, ETensorType tensor_type)
{

std::shared_ptr<void> 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<typename RByteSwap<sizeof(T)>::value_type *>(data.get()))[k] =
RByteSwap<sizeof(T)>::bswap((reinterpret_cast<const typename RByteSwap<sizeof(T)>::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<typename RByteSwap<sizeof(uint8_t)>::value_type *>(data.get()))[k] =
RByteSwap<sizeof(T)>::bswap((reinterpret_cast<const typename RByteSwap<sizeof(uint8_t)>::value_type *>(
tensorproto->raw_data().c_str()))[k]);
#endif
} else {
ExtractDataFromTP<T>::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<float>::Copy(tensorproto, data.get(), tensor_size/ 4);
break;
}
case ETensorType::DOUBLE: {
ExtractDataFromTP<double>::Copy(tensorproto, data.get(), tensor_size/ 8);
break;
}
case ETensorType::INT32: {
ExtractDataFromTP<int32_t>::Copy(tensorproto, data.get(), tensor_size/ 4);
break;
}
case ETensorType::INT64: {
ExtractDataFromTP<int64_t>::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<char *>(data.get()), buffer_size);
}

return data;
}


// Constructor of the parser
RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_unique<OperatorsMapImpl>()) {
// Register operators
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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<onnx::TensorProto *>(&graph.initializer(i));
std::vector<std::size_t> 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<ETensorType>(graph.initializer(i).data_type());
RegisterTensorType(input_name, tensor_type);

switch (tensor_type) {
case ETensorType::FLOAT: {
std::shared_ptr<void> data = GetInitializedTensorData<float>(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<void> data = GetInitializedTensorData<double>(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<void> data = GetInitializedTensorData<int32_t>(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<void> data = GetInitializedTensorData<int64_t>(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<void> data = GetInitializedTensorData<bool>(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<bool>(fLength,reinterpret_cast<bool *>(data.get(),10)) << std::endl;
break;
}
case ETensorType::UINT8: {
std::shared_ptr<void> data = GetInitializedTensorData<uint8_t>(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<uint8_t>(fLength,reinterpret_cast<uint8_t *>(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<void> 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<float*>(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<int64_t*>(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<uint8_t*>(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<bool*>(data.get())[j] << " ";
}
std::cout << std::endl;
}
}
} // end initializer list

// Initial operator order
if (verbose) {
Expand Down
Loading