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) {