From d45bce4c5aaf69e709fff17a5354dce921efa263 Mon Sep 17 00:00:00 2001 From: Jonas Rembser Date: Tue, 25 Aug 2026 08:12:32 +0000 Subject: [PATCH] [tmva][sofie] Fix reading of ONNX external weight data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The external-data reading in the ONNX parser had two problems: - The data file name was resolved once per parser instance: it defaulted to .onnx.data on the first Parse call and was never reset, and the ifstream stayed open across calls. Reusing one parser instance for several models with external data silently read every subsequent model's weights at the stored offsets of the *first* model's data file, producing garbage weights (or short reads leaving malloc'ed buffers uninitialized). - The "location" key of a tensor's external_data, which per the ONNX spec names the data file relative to the model directory, was ignored entirely, so files whose data location does not follow the .onnx.data convention could not be read at all. Resolve the data file per tensor: an explicitly set file (SetExternalDataFile) takes precedence, then the tensor's stored location relative to the model directory, then the conventional .onnx.data. Track which file is open, and reset the whole external-data state after each Parse call. This surfaced with models exported by the torch.export-based ONNX exporter of PyTorch 2.x, which stores initializers externally. Add a unit test covering the resolution of the external data file (stored location relative to the model directory, per-Parse-call default name, SetExternalDataFile precedence) and the reset of the state between Parse calls. The test hand-writes the minimal ONNX protobuf wire format, so it does not depend on the onnx Python package. 🤖 Done with the help of AI --- tmva/sofie/test/CMakeLists.txt | 5 + tmva/sofie/test/TestSofieParser.cxx | 187 ++++++++++++++++++ .../inc/TMVA/RModelParser_ONNX.hxx | 11 +- tmva/sofie_parsers/src/RModelParser_ONNX.cxx | 44 ++++- 4 files changed, 239 insertions(+), 8 deletions(-) create mode 100644 tmva/sofie/test/TestSofieParser.cxx diff --git a/tmva/sofie/test/CMakeLists.txt b/tmva/sofie/test/CMakeLists.txt index 79f21d8973439..50b574b30a7f3 100644 --- a/tmva/sofie/test/CMakeLists.txt +++ b/tmva/sofie/test/CMakeLists.txt @@ -124,6 +124,11 @@ if (BLAS_FOUND) endif() endif() +# Unit tests for the ONNX parser layer: the tests write the model files +# themselves as hand-crafted protobuf wire format, so they need no Python +ROOT_ADD_GTEST(TestSofieParser TestSofieParser.cxx + LIBRARIES Core ROOTTMVASofie ROOTTMVASofieParser) + if (ROOT_TORCH_FOUND AND ROOT_ONNX_FOUND AND NOT broken_onnx) # ModelGeneratorUtils.py holds the code shared by the generators below, so it # has to be copied next to them. diff --git a/tmva/sofie/test/TestSofieParser.cxx b/tmva/sofie/test/TestSofieParser.cxx new file mode 100644 index 0000000000000..6e952083ec312 --- /dev/null +++ b/tmva/sofie/test/TestSofieParser.cxx @@ -0,0 +1,187 @@ +// Unit tests for the SOFIE ONNX parser layer (RModelParser_ONNX). +// The model files are hand-written protobuf wire format, so the tests need +// neither the onnx Python package nor a protobuf dependency. + +#include +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include + +using namespace TMVA::Experimental::SOFIE; + +namespace { + +// --- minimal protobuf wire-format writers ---------------------------------- + +void AppendVarint(std::string &out, std::uint64_t v) +{ + while (v >= 0x80) { + out.push_back(char((v & 0x7f) | 0x80)); + v >>= 7; + } + out.push_back(char(v)); +} + +void AppendVarintField(std::string &out, int field, std::uint64_t v) +{ + AppendVarint(out, std::uint64_t(field) << 3 | 0); // wire type 0: varint + AppendVarint(out, v); +} + +void AppendBytesField(std::string &out, int field, const std::string &payload) +{ + AppendVarint(out, std::uint64_t(field) << 3 | 2); // wire type 2: length-delimited + AppendVarint(out, payload.size()); + out += payload; +} + +// StringStringEntryProto +std::string StringEntry(const std::string &key, const std::string &value) +{ + std::string out; + AppendBytesField(out, 1, key); + AppendBytesField(out, 2, value); + return out; +} + +// TensorProto for a 1-d float tensor whose data is stored externally. An +// empty location means no "location" entry, so the parser falls back to the +// conventional .data name. +std::string +ExternalFloatTensor(const std::string &name, std::uint64_t dim, std::uint64_t offset, const std::string &location = "") +{ + std::string out; + AppendVarintField(out, 1, dim); // dims + AppendVarintField(out, 2, 1); // data_type: FLOAT + AppendBytesField(out, 8, name); // name + if (!location.empty()) + AppendBytesField(out, 13, StringEntry("location", location)); // external_data + AppendBytesField(out, 13, StringEntry("offset", std::to_string(offset))); + AppendBytesField(out, 13, StringEntry("length", std::to_string(dim * sizeof(float)))); + AppendVarintField(out, 14, 1); // data_location: EXTERNAL + return out; +} + +// ModelProto holding a graph with the given initializers. The parser does +// not support graphs without nodes, so route the first tensor through an +// Identity node to the graph output. +void WriteModelFile(const std::string &fileName, const std::vector &initializers, + const std::string &firstTensorName) +{ + std::string node; + AppendBytesField(node, 1, firstTensorName); // input + AppendBytesField(node, 2, "out"); // output + AppendBytesField(node, 4, "Identity"); // op_type + + std::string output; + AppendBytesField(output, 1, "out"); // name + + std::string graph; + AppendBytesField(graph, 1, node); // node + AppendBytesField(graph, 2, "test_graph"); // name + for (const std::string &tensor : initializers) + AppendBytesField(graph, 5, tensor); // initializer + AppendBytesField(graph, 12, output); // output + + std::string model; + AppendVarintField(model, 1, 8); // ir_version + AppendBytesField(model, 7, graph); // graph + + std::ofstream file(fileName, std::ios::binary); + file.write(model.data(), model.size()); + ASSERT_TRUE(file.good()); +} + +// External weight data is little-endian per the ONNX spec, like raw_data +void WriteDataFile(const std::string &fileName, const std::vector &values, std::size_t padding = 0) +{ + std::ofstream file(fileName, std::ios::binary); + for (std::size_t i = 0; i < padding; ++i) + file.put('\0'); + for (float value : values) { + std::uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + for (int i = 0; i < 4; ++i) + file.put(char((bits >> (8 * i)) & 0xff)); + } + ASSERT_TRUE(file.good()); +} + +} // namespace + +// The "location" key of a tensor's external_data names the data file relative +// to the model directory; different tensors can point to different files. +TEST(SOFIEParser, ExternalDataLocationRelativeToModelDirectory) +{ + gSystem->mkdir("extdata_models"); + const std::vector values1{1.f, 2.f, 3.f, 4.f}; + const std::vector values2{-5.f, 6.5f}; + WriteDataFile("extdata_models/weights1.bin", values1, /*padding=*/8); + WriteDataFile("extdata_models/weights2.bin", values2); + WriteModelFile("extdata_models/modelLoc.onnx", + {ExternalFloatTensor("w1", values1.size(), /*offset=*/8, "weights1.bin"), + ExternalFloatTensor("w2", values2.size(), /*offset=*/0, "weights2.bin")}, + "w1"); + + RModelParser_ONNX parser; + RModel model = parser.Parse("extdata_models/modelLoc.onnx"); + EXPECT_EQ(model.GetTensorData("w1"), values1); + EXPECT_EQ(model.GetTensorData("w2"), values2); +} + +// A tensor without a "location" entry falls back to the conventional +// .data name, resolved per Parse call: reusing one parser +// instance for several models must not read the first model's data file. +TEST(SOFIEParser, ExternalDataFileResolvedPerParsedModel) +{ + const std::vector valuesA{10.f, 20.f, 30.f}; + const std::vector valuesB{-1.f, -2.f, -3.f}; + WriteDataFile("extdataA.onnx.data", valuesA); + WriteDataFile("extdataB.onnx.data", valuesB); + WriteModelFile("extdataA.onnx", {ExternalFloatTensor("w", valuesA.size(), /*offset=*/0)}, "w"); + WriteModelFile("extdataB.onnx", {ExternalFloatTensor("w", valuesB.size(), /*offset=*/0)}, "w"); + + RModelParser_ONNX parser; + RModel modelA = parser.Parse("extdataA.onnx"); + EXPECT_EQ(modelA.GetTensorData("w"), valuesA); + RModel modelB = parser.Parse("extdataB.onnx"); + EXPECT_EQ(modelB.GetTensorData("w"), valuesB); +} + +// A file set with SetExternalDataFile takes precedence over the stored +// location, but only for the next Parse call. +TEST(SOFIEParser, SetExternalDataFileTakesPrecedenceOnce) +{ + const std::vector valuesExplicit{5.f, 6.f}; + const std::vector valuesLocation{7.f, 8.f}; + WriteDataFile("extdata_explicit.bin", valuesExplicit); + WriteDataFile("extdata_location.bin", valuesLocation); + WriteModelFile("extdataC.onnx", + {ExternalFloatTensor("w", valuesExplicit.size(), /*offset=*/0, "extdata_location.bin")}, "w"); + + RModelParser_ONNX parser; + parser.SetExternalDataFile("extdata_explicit.bin"); + RModel modelExplicit = parser.Parse("extdataC.onnx"); + EXPECT_EQ(modelExplicit.GetTensorData("w"), valuesExplicit); + + RModel modelLocation = parser.Parse("extdataC.onnx"); + EXPECT_EQ(modelLocation.GetTensorData("w"), valuesLocation); +} + +// A data file that cannot be opened is an error, not silently zeroed weights +TEST(SOFIEParser, MissingExternalDataFileThrows) +{ + WriteModelFile("extdataD.onnx", {ExternalFloatTensor("w", 2, /*offset=*/0, "extdata_does_not_exist.bin")}, "w"); + + RModelParser_ONNX parser; + EXPECT_THROW(parser.Parse("extdataD.onnx"), std::runtime_error); +} diff --git a/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx b/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx index fe7365006c942..11c128c80511b 100644 --- a/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx +++ b/tmva/sofie_parsers/inc/TMVA/RModelParser_ONNX.hxx @@ -46,8 +46,14 @@ private: // weight data file std::ifstream fDataFile; - // filename of model + // user-provided external data file name (see SetExternalDataFile), valid for the next Parse call std::string fDataFileName; + // directory of the model being parsed, used to resolve relative external data locations + std::string fModelDirectory; + // default external data file name (.data), used when a tensor provides no location + std::string fDefaultDataFileName; + // name of the external data file fDataFile currently has open + std::string fOpenedDataFileName; public: @@ -89,6 +95,9 @@ public: std::shared_ptr GetInitializedTensorData(onnx::TensorProto *tensorproto, size_t tensor_length, ETensorType type ); + // reset the external-data reading state after parsing a model + void ResetExternalDataState(); + public: RModelParser_ONNX() noexcept; diff --git a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx index 099fb62c0fd21..669d9e57f4024 100644 --- a/tmva/sofie_parsers/src/RModelParser_ONNX.cxx +++ b/tmva/sofie_parsers/src/RModelParser_ONNX.cxx @@ -253,8 +253,6 @@ std::shared_ptr RModelParser_ONNX::GetInitializedTensorData(onnx::TensorPr } 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; @@ -265,17 +263,33 @@ std::shared_ptr RModelParser_ONNX::GetInitializedTensorData(onnx::TensorPr else if (kv.key() == "offset") offset = std::stoull(kv.value()); else if (kv.key() == "length") buffer_size = std::stoull(kv.value()); } + + // an explicitly set data file (SetExternalDataFile) takes precedence; + // otherwise use the location stored in the model, which is a path + // relative to the model directory, and as a last resort the + // conventional .data + std::string dataFileName = fDataFileName; + if (dataFileName.empty()) + dataFileName = location.empty() ? fDefaultDataFileName : fModelDirectory + location; + if (dataFileName.empty()) + throw std::runtime_error("TMVA::SOFIE ONNX : tensor " + tensorproto->name() + + " has external data but no data file location is available"); + if (fVerbose) - std::cout << " at location " << location << " offset " << offset << " and with length " << buffer_size << std::endl; + std::cout << "Initialized data are stored externally in file " << dataFileName + << " 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 + // open the data file if needed (a previous tensor may have opened a different one) + if (fDataFile.is_open() && fOpenedDataFileName != dataFileName) + fDataFile.close(); if (!fDataFile.is_open()) { - fDataFile.open(fDataFileName, std::ios::binary); + fDataFile.open(dataFileName, std::ios::binary); if (!fDataFile.is_open()) - throw std::runtime_error("TMVA::SOFIE ONNX: error reading external weight ONNX data file " + fDataFileName); + throw std::runtime_error("TMVA::SOFIE ONNX: error reading external weight ONNX data file " + dataFileName); + fOpenedDataFileName = dataFileName; } fDataFile.seekg(offset); @@ -539,10 +553,12 @@ 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"; + fModelDirectory = (isep != std::string::npos) ? filename.substr(0, isep + 1) : ""; + fDefaultDataFileName = filename + ".data"; RModel rmodel(filename_nodir, parsetime); ParseONNXGraph(rmodel, graph, filename_nodir); + ResetExternalDataState(); return rmodel; } @@ -564,9 +580,23 @@ RModel RModelParser_ONNX::Parse(std::istream &input, std::string const &name, bo RModel rmodel(name, parsetime); ParseONNXGraph(rmodel, graph, name); + ResetExternalDataState(); return rmodel; } +// Reset the state used to read external weight data, so that the next Parse +// call does not pick up the data file of a previously parsed model. The +// file name set with SetExternalDataFile is valid for a single Parse call. +void RModelParser_ONNX::ResetExternalDataState() +{ + fDataFileName.clear(); + fModelDirectory.clear(); + fDefaultDataFileName.clear(); + fOpenedDataFileName.clear(); + if (fDataFile.is_open()) + fDataFile.close(); +} + std::unique_ptr RModelParser_ONNX::LoadModel(const std::string &filename) { std::fstream input(filename, std::ios::in | std::ios::binary); if (!input) {