diff --git a/Common/ML/include/ML/OrtInterface.h b/Common/ML/include/ML/OrtInterface.h index 04a5e0ba5c9fc..987ce8fb4d6dd 100644 --- a/Common/ML/include/ML/OrtInterface.h +++ b/Common/ML/include/ML/OrtInterface.h @@ -51,6 +51,7 @@ class OrtModel void initOptions(std::unordered_map optionsMap); void initEnvironment(); void initSession(); + void initSessionFromBuffer(const char* buffer, size_t bufferSize); void memoryOnDevice(int32_t = 0); bool isInitialized() { return mInitialized; } void resetSession(); diff --git a/Common/ML/src/OrtInterface.cxx b/Common/ML/src/OrtInterface.cxx index d30d05d1d1a00..8f88ab18dacbd 100644 --- a/Common/ML/src/OrtInterface.cxx +++ b/Common/ML/src/OrtInterface.cxx @@ -138,6 +138,24 @@ void OrtModel::initEnvironment() (mPImplOrt->env)->DisableTelemetryEvents(); // Disable telemetry events } +void OrtModel::initSessionFromBuffer(const char* buffer, size_t bufferSize) +{ + mPImplOrt->sessionOptions.AddConfigEntry("session.load_model_format", "ONNX"); + mPImplOrt->sessionOptions.AddConfigEntry("session.use_ort_model_bytes_directly", "1"); + + mPImplOrt->session = std::make_unique(*mPImplOrt->env, + buffer, + bufferSize, + mPImplOrt->sessionOptions); + mPImplOrt->ioBinding = std::make_unique(*mPImplOrt->session); + + setIO(); + + if (mLoggingLevel < 2) { + LOG(info) << "(ORT) Model loaded successfully from buffer! (inputs: " << printShape(mInputShapes, mInputNames) << ", outputs: " << printShape(mOutputShapes, mInputNames) << ")"; + } +} + void OrtModel::initSession() { if (mAllocateDeviceMemory) { diff --git a/Detectors/TPC/base/test/testTPCCDBInterface.cxx b/Detectors/TPC/base/test/testTPCCDBInterface.cxx index 3074c5e90a00c..5a5384a4134ed 100644 --- a/Detectors/TPC/base/test/testTPCCDBInterface.cxx +++ b/Detectors/TPC/base/test/testTPCCDBInterface.cxx @@ -22,7 +22,6 @@ // o2 includes #include "TPCBase/CDBInterface.h" -#include "TPCBase/CDBInterface.h" #include "TPCBase/CalArray.h" #include "TPCBase/CalDet.h" #include "TPCBase/Mapper.h" diff --git a/Detectors/TPC/calibration/CMakeLists.txt b/Detectors/TPC/calibration/CMakeLists.txt index 8bcb3254edb32..e5cc25230d2fc 100644 --- a/Detectors/TPC/calibration/CMakeLists.txt +++ b/Detectors/TPC/calibration/CMakeLists.txt @@ -25,7 +25,6 @@ o2_add_library(TPCCalibration src/CalibPadGainTracksBase.cxx src/CalibLaserTracks.cxx src/LaserTracksCalibrator.cxx - src/NeuralNetworkClusterizer.cxx src/SACDecoder.cxx src/IDCAverageGroup.cxx src/IDCAverageGroupBase.cxx @@ -84,7 +83,6 @@ o2_target_root_dictionary(TPCCalibration include/TPCCalibration/FastHisto.h include/TPCCalibration/CalibLaserTracks.h include/TPCCalibration/LaserTracksCalibrator.h - include/TPCCalibration/NeuralNetworkClusterizer.h include/TPCCalibration/SACDecoder.h include/TPCCalibration/IDCAverageGroup.h include/TPCCalibration/IDCAverageGroupBase.h diff --git a/Detectors/TPC/calibration/include/TPCCalibration/NeuralNetworkClusterizer.h b/Detectors/TPC/calibration/include/TPCCalibration/NeuralNetworkClusterizer.h deleted file mode 100644 index 196bba644714c..0000000000000 --- a/Detectors/TPC/calibration/include/TPCCalibration/NeuralNetworkClusterizer.h +++ /dev/null @@ -1,38 +0,0 @@ -// Copyright 2019-2020 CERN and copyright holders of ALICE O2. -// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. -// All rights not expressly granted are reserved. -// -// This software is distributed under the terms of the GNU General Public -// License v3 (GPL Version 3), copied verbatim in the file "COPYING". -// -// In applying this license CERN does not waive the privileges and immunities -// granted to it by virtue of its status as an Intergovernmental Organization -// or submit itself to any jurisdiction. - -/// \file NeuralNetworkClusterizer.h -/// \brief Fetching neural networks for clusterization from CCDB -/// \author Christian Sonnabend - -#ifndef AliceO2_TPC_NeuralNetworkClusterizer_h -#define AliceO2_TPC_NeuralNetworkClusterizer_h - -#include "CCDB/CcdbApi.h" - -namespace o2::tpc -{ - -class NeuralNetworkClusterizer -{ - public: - NeuralNetworkClusterizer() = default; - void initCcdbApi(std::string url); - void loadIndividualFromCCDB(std::map settings); - - private: - o2::ccdb::CcdbApi ccdbApi; - std::map metadata; - std::map headers; -}; - -} // namespace o2::tpc -#endif diff --git a/Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx b/Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx deleted file mode 100644 index bfbb7afc946f8..0000000000000 --- a/Detectors/TPC/calibration/src/NeuralNetworkClusterizer.cxx +++ /dev/null @@ -1,48 +0,0 @@ -// Copyright 2019-2020 CERN and copyright holders of ALICE O2. -// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. -// All rights not expressly granted are reserved. -// -// This software is distributed under the terms of the GNU General Public -// License v3 (GPL Version 3), copied verbatim in the file "COPYING". -// -// In applying this license CERN does not waive the privileges and immunities -// granted to it by virtue of its status as an Intergovernmental Organization -// or submit itself to any jurisdiction. - -/// \file NeuralNetworkClusterizer.cxx -/// \brief Fetching neural networks for clusterization from CCDB -/// \author Christian Sonnabend - -#include -#include "TPCCalibration/NeuralNetworkClusterizer.h" - -using namespace o2::tpc; - -void NeuralNetworkClusterizer::initCcdbApi(std::string url) -{ - ccdbApi.init(url); -} - -void NeuralNetworkClusterizer::loadIndividualFromCCDB(std::map settings) -{ - metadata["inputDType"] = settings["inputDType"]; - metadata["outputDType"] = settings["outputDType"]; - metadata["nnCCDBEvalType"] = settings["nnCCDBEvalType"]; // classification_1C, classification_2C, regression_1C, regression_2C - metadata["nnCCDBWithMomentum"] = settings["nnCCDBWithMomentum"]; // 0, 1 -> Only for regression model - metadata["nnCCDBLayerType"] = settings["nnCCDBLayerType"]; // FC, CNN - if (settings["nnCCDBInteractionRate"] != "" && std::stoi(settings["nnCCDBInteractionRate"]) > 0) { - metadata["nnCCDBInteractionRate"] = settings["nnCCDBInteractionRate"]; - } - if (settings["nnCCDBBeamType"] != "") { - metadata["nnCCDBBeamType"] = settings["nnCCDBBeamType"]; - } - - bool retrieveSuccess = ccdbApi.retrieveBlob(settings["nnCCDBPath"], settings["outputFolder"], metadata, 1, false, settings["outputFile"]); - // headers = ccdbApi.retrieveHeaders(settings["nnPathCCDB"], metadata, 1); // potentially needed to init some local variables - - if (retrieveSuccess) { - LOG(info) << "Network " << settings["nnCCDBPath"] << " retrieved from CCDB, stored at " << settings["outputFile"]; - } else { - LOG(error) << "Failed to retrieve network from CCDB"; - } -} diff --git a/GPU/GPUTracking/CMakeLists.txt b/GPU/GPUTracking/CMakeLists.txt index 2a0b9b9edfa09..6dd718f07a9f1 100644 --- a/GPU/GPUTracking/CMakeLists.txt +++ b/GPU/GPUTracking/CMakeLists.txt @@ -209,6 +209,7 @@ set(SRCS_DATATYPES DataTypes/TPCPadBitMap.cxx DataTypes/TPCZSLinkMapping.cxx DataTypes/CalibdEdxContainer.cxx + DataTypes/ORTRootSerializer.cxx DataTypes/CalibdEdxTrackTopologyPol.cxx DataTypes/CalibdEdxTrackTopologySpline.cxx DataTypes/GPUTRDTrackO2.cxx) diff --git a/GPU/GPUTracking/DataTypes/GPUDataTypes.h b/GPU/GPUTracking/DataTypes/GPUDataTypes.h index 967d6a73914dd..8bf8084e048fd 100644 --- a/GPU/GPUTracking/DataTypes/GPUDataTypes.h +++ b/GPU/GPUTracking/DataTypes/GPUDataTypes.h @@ -85,6 +85,7 @@ class Cluster; namespace tpc { class CalibdEdxContainer; +class ORTRootSerializer; } // namespace tpc } // namespace o2 @@ -182,6 +183,9 @@ struct GPUCalibObjectsTemplate { // use only pointers on PODs or flat objects he typename S::type* dEdxCalibContainer = nullptr; typename S>::type* o2Propagator = nullptr; typename S::type* itsPatternDict = nullptr; + + // NN clusterizer objects + typename S::type* nnClusterizerNetworks[3] = {nullptr, nullptr, nullptr}; }; typedef GPUCalibObjectsTemplate GPUCalibObjects; // NOTE: These 2 must have identical layout since they are memcopied typedef GPUCalibObjectsTemplate GPUCalibObjectsConst; diff --git a/GPU/GPUTracking/DataTypes/ORTRootSerializer.cxx b/GPU/GPUTracking/DataTypes/ORTRootSerializer.cxx new file mode 100644 index 0000000000000..82a8be1fdfec8 --- /dev/null +++ b/GPU/GPUTracking/DataTypes/ORTRootSerializer.cxx @@ -0,0 +1,25 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file ORTRootSerializer.cxx +/// \author Christian Sonnabend + +#include "ORTRootSerializer.h" +#include + +using namespace o2::tpc; + +/// Initialize the serialization from a char* buffer containing the model +void ORTRootSerializer::setOnnxModel(const char* onnxModel, uint32_t size) +{ + mModelBuffer.resize(size); + std::memcpy(mModelBuffer.data(), onnxModel, size); +} diff --git a/GPU/GPUTracking/DataTypes/ORTRootSerializer.h b/GPU/GPUTracking/DataTypes/ORTRootSerializer.h new file mode 100644 index 0000000000000..24009d4435a96 --- /dev/null +++ b/GPU/GPUTracking/DataTypes/ORTRootSerializer.h @@ -0,0 +1,43 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file ORTRootSerializer.h +/// \brief Class to serialize ONNX objects for ROOT snapshots of CCDB objects at runtime +/// \author Christian Sonnabend + +#ifndef ALICEO2_TPC_ORTROOTSERIALIZER_H_ +#define ALICEO2_TPC_ORTROOTSERIALIZER_H_ + +#include "GPUCommonRtypes.h" +#include +#include + +namespace o2::tpc +{ + +class ORTRootSerializer +{ + public: + ORTRootSerializer() = default; + ~ORTRootSerializer() = default; + + void setOnnxModel(const char* onnxModel, uint32_t size); + const char* getONNXModel() const { return mModelBuffer.data(); } + uint32_t getONNXModelSize() const { return static_cast(mModelBuffer.size()); } + + private: + std::vector mModelBuffer; ///< buffer for serialization + ClassDefNV(ORTRootSerializer, 1); +}; + +} // namespace o2::tpc + +#endif // ALICEO2_TPC_ORTROOTSERIALIZER_H_ diff --git a/GPU/GPUTracking/Definitions/GPUSettingsList.h b/GPU/GPUTracking/Definitions/GPUSettingsList.h index 052da8ae54c60..dc1742453ef39 100644 --- a/GPU/GPUTracking/Definitions/GPUSettingsList.h +++ b/GPU/GPUTracking/Definitions/GPUSettingsList.h @@ -277,22 +277,22 @@ AddOption(nnClusterizerBoundaryFillValue, int, -1, "", 0, "Fill value for the bo AddOption(nnClusterizerApplyNoiseSuppression, int, 1, "", 0, "Applies the NoiseSuppression kernel before the digits to the network are filled") AddOption(nnClusterizerSetDeconvolutionFlags, int, 1, "", 0, "Runs the deconvolution kernel without overwriting the charge in order to make cluster-to-track attachment identical to heuristic CF") AddOption(nnClassificationPath, std::string, "network_class.onnx", "", 0, "The classification network path") -AddOption(nnClassThreshold, float, 0.5, "", 0, "The cutoff at which clusters will be accepted / rejected.") AddOption(nnRegressionPath, std::string, "network_reg.onnx", "", 0, "The regression network path") +AddOption(nnClassThreshold, float, 0.5, "", 0, "The cutoff at which clusters will be accepted / rejected.") AddOption(nnSigmoidTrafoClassThreshold, int, 1, "", 0, "If true (default), then the classification threshold is transformed by an inverse sigmoid function. This depends on how the network was trained (with a sigmoid as acitvation function in the last layer or not).") AddOption(nnEvalMode, std::string, "c1:r1", "", 0, "Concatention of modes, e.g. c1:r1 (classification class 1, regression class 1)") AddOption(nnClusterizerUseClassification, int, 1, "", 0, "If 1, the classification output of the network is used to select clusters, else only the regression output is used and no clusters are rejected by classification") AddOption(nnClusterizerForceGpuInputFill, int, 0, "", 0, "Forces to use the fillInputNNGPU function") // CCDB AddOption(nnLoadFromCCDB, int, 0, "", 0, "If 1 networks are fetched from ccdb, else locally") +AddOption(nnCCDBDumpToFile, int, 0, "", 0, "If 1, additionally dump fetched CCDB networks to nnLocalFolder") AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched") -AddOption(nnCCDBURL, std::string, "http://ccdb-test.cern.ch:8080", "", 0, "The CCDB URL from where the network files are fetched") AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks") -AddOption(nnCCDBWithMomentum, int, 1, "", 0, "Distinguishes between the network with and without momentum output for the regression") +AddOption(nnCCDBWithMomentum, std::string, "", "", 0, "Distinguishes between the network with and without momentum output for the regression") AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN") -AddOption(nnCCDBRegressionLayerType, std::string, "CNN", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN") -AddOption(nnCCDBBeamType, std::string, "PbPb", "", 0, "Distinguishes between networks trained for different beam types. Options: PbPb, pp") -AddOption(nnCCDBInteractionRate, int, 50, "", 0, "Distinguishes between networks for different interaction rates [kHz].") +AddOption(nnCCDBRegressionLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN") +AddOption(nnCCDBBeamType, std::string, "pp", "", 0, "Distinguishes between networks trained for different beam types. Options: pp, pPb, PbPb") +AddOption(nnCCDBInteractionRate, std::string, "500", "", 0, "Distinguishes between networks for different interaction rates [kHz].") AddHelp("help", 'h') EndConfig() diff --git a/GPU/GPUTracking/GPUTrackingLinkDef_O2_DataTypes.h b/GPU/GPUTracking/GPUTrackingLinkDef_O2_DataTypes.h index 46fd50464c69b..7bd2c689c5354 100644 --- a/GPU/GPUTracking/GPUTrackingLinkDef_O2_DataTypes.h +++ b/GPU/GPUTracking/GPUTrackingLinkDef_O2_DataTypes.h @@ -43,5 +43,6 @@ #pragma link C++ class o2::tpc::CalibdEdxTrackTopologyPol + ; #pragma link C++ class o2::tpc::CalibdEdxTrackTopologySpline + ; #pragma link C++ struct o2::tpc::CalibdEdxTrackTopologyPolContainer + ; +#pragma link C++ struct o2::tpc::ORTRootSerializer + ; #endif diff --git a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx index bfb0457744ce5..5426f0eafdad6 100644 --- a/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx +++ b/GPU/GPUTracking/Global/GPUChainTrackingClusterizer.cxx @@ -47,6 +47,7 @@ #ifdef GPUCA_HAS_ONNX #include "GPUTPCNNClusterizerKernels.h" #include "GPUTPCNNClusterizerHost.h" +#include "ORTRootSerializer.h" #endif #ifdef GPUCA_O2_LIB @@ -639,7 +640,7 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput) // Maximum of 4 lanes supported HighResTimer* nnTimers[12]; - if (GetProcessingSettings().nn.applyNNclusterizer) { + if (nn_settings.applyNNclusterizer) { int32_t deviceId = -1; int32_t numLanes = GetProcessingSettings().nTPCClustererLanes; int32_t maxThreads = mRec->getNKernelHostThreads(true); @@ -677,7 +678,11 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput) // nnApplications[lane].directOrtAllocator((nnApplications[lane].mModelClass).getEnv(), (nnApplications[lane].mModelClass).getMemoryInfo(), mRec, recreateMemoryAllocator); // } // recreateMemoryAllocator = true; - (nnApplications[lane].mModelClass).initSession(); + if (!nn_settings.nnLoadFromCCDB) { + (nnApplications[lane].mModelClass).initSession(); // loads from file + } else { + (nnApplications[lane].mModelClass).initSessionFromBuffer((processors()->calibObjects.nnClusterizerNetworks[0])->getONNXModel(), (processors()->calibObjects.nnClusterizerNetworks[0])->getONNXModelSize()); // loads from CCDB + } } if (nnApplications[lane].mModelsUsed[1]) { SetONNXGPUStream(*(nnApplications[lane].mModelReg1).getSessionOptions(), lane, &deviceId); @@ -688,7 +693,11 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput) // (nnApplications[lane].mModelReg1).setEnv((nnApplications[lane].mModelClass).getEnv()); (nnApplications[lane].mModelReg1).initEnvironment(); // nnApplications[lane].directOrtAllocator((nnApplications[lane].mModelReg1).getEnv(), (nnApplications[lane].mModelReg1).getMemoryInfo(), mRec, recreateMemoryAllocator); - (nnApplications[lane].mModelReg1).initSession(); + if (!nn_settings.nnLoadFromCCDB) { + (nnApplications[lane].mModelReg1).initSession(); // loads from file + } else { + (nnApplications[lane].mModelReg1).initSessionFromBuffer((processors()->calibObjects.nnClusterizerNetworks[1])->getONNXModel(), (processors()->calibObjects.nnClusterizerNetworks[1])->getONNXModelSize()); // loads from CCDB + } } if (nnApplications[lane].mModelsUsed[2]) { SetONNXGPUStream(*(nnApplications[lane].mModelReg2).getSessionOptions(), lane, &deviceId); @@ -699,7 +708,11 @@ int32_t GPUChainTracking::RunTPCClusterizer(bool synchronizeOutput) // (nnApplications[lane].mModelReg2).setEnv((nnApplications[lane].mModelClass).getEnv()); (nnApplications[lane].mModelReg2).initEnvironment(); // nnApplications[lane].directOrtAllocator((nnApplications[lane].mModelClass).getEnv(), (nnApplications[lane].mModelClass).getMemoryInfo(), mRec, recreateMemoryAllocator); - (nnApplications[lane].mModelReg2).initSession(); + if (!nn_settings.nnLoadFromCCDB) { + (nnApplications[lane].mModelReg2).initSession(); // loads from file + } else { + (nnApplications[lane].mModelReg2).initSessionFromBuffer((processors()->calibObjects.nnClusterizerNetworks[2])->getONNXModel(), (processors()->calibObjects.nnClusterizerNetworks[2])->getONNXModelSize()); // loads from CCDB + } } if (nn_settings.nnClusterizerVerbosity > 0) { LOG(info) << "(ORT) Allocated ONNX stream for lane " << lane << " and device " << deviceId; diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx index 582a0c6d7435a..77d5ee13f85fb 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx +++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx @@ -36,7 +36,7 @@ void GPUTPCNNClusterizerHost::init(const GPUSettingsProcessingNNclusterizer& set std::vector evalMode = o2::utils::Str::tokenize(settings.nnEvalMode, ':'); if (settings.nnLoadFromCCDB) { - reg_model_path = settings.nnLocalFolder + "/net_regression_c1.onnx"; // Needs to be set identical to NeuralNetworkClusterizer.cxx, otherwise the networks might be loaded from the wrong place + reg_model_path = settings.nnLocalFolder + "/net_regression_c1.onnx"; // Needs to be set identical to GPUWorkflowSpec.cxx, otherwise the networks might be loaded from the wrong place if (evalMode[0] == "c1") { class_model_path = settings.nnLocalFolder + "/net_classification_c1.onnx"; } else if (evalMode[0] == "c2") { diff --git a/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h b/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h index 160efd4048af0..d610269abca81 100644 --- a/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h +++ b/GPU/Workflow/include/GPUWorkflow/GPUWorkflowSpec.h @@ -135,6 +135,11 @@ class GPURecoWorkflowSpec : public o2::framework::Task bool tpcTriggerHandling = false; bool isITS3 = false; bool useFilteredOutputSpecs = false; + + // NN clusterizer + bool nnLoadFromCCDB = false; + bool nnDumpToFile = false; + std::vector nnEvalMode; }; GPURecoWorkflowSpec(CompletionPolicyData* policyData, Config const& specconfig, std::vector const& tpcsectors, uint64_t tpcSectorMask, std::shared_ptr& ggr, std::function** gPolicyOrder = nullptr); @@ -230,7 +235,7 @@ class GPURecoWorkflowSpec : public o2::framework::Task uint32_t mNextThreadIndex = 0; bool mUpdateGainMapCCDB = true; std::unique_ptr mTFSettings; - std::unique_ptr mNNClusterizerSettings; + std::map nnCCDBSettings; Config mSpecConfig; std::shared_ptr mGGR; diff --git a/GPU/Workflow/src/GPUWorkflowSpec.cxx b/GPU/Workflow/src/GPUWorkflowSpec.cxx index d3d3eb14869e0..d7ea772c31653 100644 --- a/GPU/Workflow/src/GPUWorkflowSpec.cxx +++ b/GPU/Workflow/src/GPUWorkflowSpec.cxx @@ -54,6 +54,7 @@ #include "GPUO2Interface.h" #include "GPUO2InterfaceUtils.h" #include "CalibdEdxContainer.h" +#include "ORTRootSerializer.h" #include "GPUNewCalibValues.h" #include "TPCPadGainCalib.h" #include "TPCZSLinkMapping.h" @@ -78,7 +79,6 @@ #include "DetectorsRaw/RDHUtils.h" #include "ITStracking/TrackingInterface.h" #include "GPUWorkflowInternal.h" -#include "TPCCalibration/NeuralNetworkClusterizer.h" // #include "Framework/ThreadPool.h" #include @@ -133,50 +133,6 @@ void GPURecoWorkflowSpec::init(InitContext& ic) { GRPGeomHelper::instance().setRequest(mGGR); GPUO2InterfaceConfiguration& config = *mConfig.get(); - GPUSettingsProcessingNNclusterizer& mNNClusterizerSettings = mConfig->configProcessing.nn; - - if (mNNClusterizerSettings.nnLoadFromCCDB) { - LOG(info) << "Loading neural networks from CCDB"; - o2::tpc::NeuralNetworkClusterizer nnClusterizerFetcher; - nnClusterizerFetcher.initCcdbApi(mNNClusterizerSettings.nnCCDBURL); - std::map ccdbSettings = { - {"nnCCDBURL", mNNClusterizerSettings.nnCCDBURL}, - {"nnCCDBPath", mNNClusterizerSettings.nnCCDBPath}, - {"inputDType", mNNClusterizerSettings.nnInferenceInputDType}, - {"outputDType", mNNClusterizerSettings.nnInferenceOutputDType}, - {"outputFolder", mNNClusterizerSettings.nnLocalFolder}, - {"nnCCDBPath", mNNClusterizerSettings.nnCCDBPath}, - {"nnCCDBWithMomentum", std::to_string(mNNClusterizerSettings.nnCCDBWithMomentum)}, - {"nnCCDBBeamType", mNNClusterizerSettings.nnCCDBBeamType}, - {"nnCCDBInteractionRate", std::to_string(mNNClusterizerSettings.nnCCDBInteractionRate)}}; - - std::string nnFetchFolder = mNNClusterizerSettings.nnLocalFolder; - std::vector evalMode = o2::utils::Str::tokenize(mNNClusterizerSettings.nnEvalMode, ':'); - - if (evalMode[0] == "c1") { - ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBClassificationLayerType; - ccdbSettings["nnCCDBEvalType"] = "classification_c1"; - ccdbSettings["outputFile"] = "net_classification_c1.onnx"; - nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings); - } else if (evalMode[0] == "c2") { - ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBClassificationLayerType; - ccdbSettings["nnCCDBEvalType"] = "classification_c2"; - ccdbSettings["outputFile"] = "net_classification_c2.onnx"; - nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings); - } - - ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBRegressionLayerType; - ccdbSettings["nnCCDBEvalType"] = "regression_c1"; - ccdbSettings["outputFile"] = "net_regression_c1.onnx"; - nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings); - if (evalMode[1] == "r2") { - ccdbSettings["nnCCDBLayerType"] = mNNClusterizerSettings.nnCCDBRegressionLayerType; - ccdbSettings["nnCCDBEvalType"] = "regression_c2"; - ccdbSettings["outputFile"] = "net_regression_c2.onnx"; - nnClusterizerFetcher.loadIndividualFromCCDB(ccdbSettings); - } - LOG(info) << "Neural network loading done!"; - } // Create configuration object and fill settings mConfig->configGRP.solenoidBzNominalGPU = 0; @@ -185,6 +141,7 @@ void GPURecoWorkflowSpec::init(InitContext& ic) mTFSettings->simStartOrbit = hbfu.getFirstIRofTF(o2::InteractionRecord(0, hbfu.orbitFirstSampled)).orbit; *mConfParam = mConfig->ReadConfigurableParam(); + if (mConfParam->display) { mDisplayFrontend.reset(GPUDisplayFrontendInterface::getFrontend(mConfig->configDisplay.displayFrontend.c_str())); mConfig->configProcessing.eventDisplay = mDisplayFrontend.get(); @@ -1124,6 +1081,27 @@ void GPURecoWorkflowSpec::doCalibUpdates(o2::framework::ProcessingContext& pc, c newCalibValues.tpcTimeBinCut = mConfig->configGRP.tpcCutTimeBin = mTPCCutAtTimeBin; needCalibUpdate = true; } + if (mSpecConfig.nnLoadFromCCDB) { + auto dumpToFile = [](const char* buffer, std::size_t validSize, const std::string& path) { + std::ofstream out(path, std::ios::binary | std::ios::trunc); + if (!out.is_open()) { + throw std::runtime_error("Failed to open output file: " + path); + } + + out.write(buffer, static_cast(validSize)); + if (!out) { + throw std::runtime_error("Failed while writing data to: " + path); + } + }; + for (int i = 0; i < 3; i++) { + newCalibObjects.nnClusterizerNetworks[i] = mConfig->configCalib.nnClusterizerNetworks[i]; + if (mSpecConfig.nnDumpToFile && newCalibObjects.nnClusterizerNetworks[i]) { + std::string path = "tpc_nn_clusterizer_" + std::to_string(i) + ".onnx"; + dumpToFile(newCalibObjects.nnClusterizerNetworks[i]->getONNXModel(), newCalibObjects.nnClusterizerNetworks[i]->getONNXModelSize(), path); + LOG(info) << "Dumped TPC clusterizer NN " << i << " to file " << path; + } + } + } if (needCalibUpdate) { LOG(info) << "Updating GPUReconstruction calibration objects"; mGPUReco->UpdateCalibration(newCalibObjects, newCalibValues); @@ -1262,6 +1240,67 @@ Inputs GPURecoWorkflowSpec::inputs() } } + // NN clusterizer + *mConfParam = mConfig->ReadConfigurableParam(); + if (mConfig->configProcessing.nn.nnLoadFromCCDB) { + + LOG(info) << "(NN CLUS) Enabling fetching of TPC NN clusterizer from CCDB"; + mSpecConfig.nnLoadFromCCDB = true; + mSpecConfig.nnDumpToFile = mConfig->configProcessing.nn.nnCCDBDumpToFile; + GPUSettingsProcessingNNclusterizer& nnClusterizerSettings = mConfig->configProcessing.nn; + + std::map metadata; + metadata["inputDType"] = nnClusterizerSettings.nnInferenceInputDType; // FP16 or FP32 + metadata["outputDType"] = nnClusterizerSettings.nnInferenceOutputDType; // FP16 or FP32 + metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model + metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBClassificationLayerType; // FC, CNN + metadata["nnCCDBInteractionRate"] = nnClusterizerSettings.nnCCDBInteractionRate; // in kHz + metadata["nnCCDBBeamType"] = nnClusterizerSettings.nnCCDBBeamType; // pp, pPb, PbPb + + auto convert_map_to_metadata = [](const std::map& inputMap, std::vector& outputMetadata) { + for (const auto& [key, value] : inputMap) { + if (value != "") { + outputMetadata.push_back({key, value}); + } + } + }; + + mSpecConfig.nnEvalMode = o2::utils::Str::tokenize(nnClusterizerSettings.nnEvalMode, ':'); + std::vector ccdb_metadata; + + if (mConfParam->printSettings) { + auto printSettings = [](const std::map& settings) { + LOG(info) << "(NN CLUS) NN Clusterizer CCDB settings:"; + for (const auto& [key, value] : settings) { + LOG(info) << " " << key << " : " << value; + } + }; + printSettings(metadata); + } + + if (mSpecConfig.nnEvalMode[0] == "c1") { + metadata["nnCCDBEvalType"] = "classification_c1"; + convert_map_to_metadata(metadata, ccdb_metadata); + inputs.emplace_back("nn_classification_c1", gDataOriginTPC, "NNCLUSTERIZER_C1", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); + } else if (mSpecConfig.nnEvalMode[0] == "c2") { + metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType; + metadata["nnCCDBEvalType"] = "classification_c2"; + convert_map_to_metadata(metadata, ccdb_metadata); + inputs.emplace_back("nn_classification_c2", gDataOriginTPC, "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); + } + + metadata["nnCCDBEvalType"] = "regression_c1"; + metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType; + convert_map_to_metadata(metadata, ccdb_metadata); + inputs.emplace_back("nn_regression_c1", gDataOriginTPC, "NNCLUSTERIZER_R1", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); + + if (mSpecConfig.nnEvalMode[1] == "r2") { + metadata["nnCCDBEvalType"] = "regression_c2"; + convert_map_to_metadata(metadata, ccdb_metadata); + inputs.emplace_back("nn_regression_c2", gDataOriginTPC, "NNCLUSTERIZER_R2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); + } + } + return inputs; }; diff --git a/GPU/Workflow/src/GPUWorkflowTPC.cxx b/GPU/Workflow/src/GPUWorkflowTPC.cxx index 6606386819b64..13a3c4b6162b8 100644 --- a/GPU/Workflow/src/GPUWorkflowTPC.cxx +++ b/GPU/Workflow/src/GPUWorkflowTPC.cxx @@ -49,6 +49,7 @@ #include "GPUO2Interface.h" #include "GPUO2InterfaceUtils.h" #include "CalibdEdxContainer.h" +#include "ORTRootSerializer.h" #include "GPUNewCalibValues.h" #include "TPCPadGainCalib.h" #include "TPCZSLinkMapping.h" @@ -293,6 +294,18 @@ void GPURecoWorkflowSpec::finaliseCCDBTPC(ConcreteDataMatcher& matcher, void* ob mTPCDeadChannelMapCreator->getDeadChannelMapFEE().getSum(), mTPCDeadChannelMapCreator->getDeadChannelMap().getSum()); } else if (mTPCVDriftHelper->accountCCDBInputs(matcher, obj)) { } else if (mCalibObjects.mFastTransformHelper->accountCCDBInputs(matcher, obj)) { + } else if (matcher == ConcreteDataMatcher(gDataOriginTPC, "NNCLUSTERIZER_C1", 0)) { + mConfig->configCalib.nnClusterizerNetworks[0] = static_cast(obj); + LOG(info) << "(NN CLUS) " << (mConfig->configCalib.nnClusterizerNetworks[0])->getONNXModelSize() << " bytes loaded for NN clusterizer: classification_c1"; + } else if (matcher == ConcreteDataMatcher(gDataOriginTPC, "NNCLUSTERIZER_C2", 0)) { + mConfig->configCalib.nnClusterizerNetworks[0] = static_cast(obj); + LOG(info) << "(NN CLUS) " << (mConfig->configCalib.nnClusterizerNetworks[0])->getONNXModelSize() << " bytes loaded for NN clusterizer: classification_c2"; + } else if (matcher == ConcreteDataMatcher(gDataOriginTPC, "NNCLUSTERIZER_R1", 0)) { + mConfig->configCalib.nnClusterizerNetworks[1] = static_cast(obj); + LOG(info) << "(NN CLUS) " << (mConfig->configCalib.nnClusterizerNetworks[1])->getONNXModelSize() << " bytes loaded for NN clusterizer: regression_c1"; + } else if (matcher == ConcreteDataMatcher(gDataOriginTPC, "NNCLUSTERIZER_R2", 0)) { + mConfig->configCalib.nnClusterizerNetworks[2] = static_cast(obj); + LOG(info) << "(NN CLUS) " << (mConfig->configCalib.nnClusterizerNetworks[2])->getONNXModelSize() << " bytes loaded for NN clusterizer: regression_c2"; } } @@ -405,6 +418,21 @@ bool GPURecoWorkflowSpec::fetchCalibsCCDBTPC(ProcessingCon newCalibObjects.tpcPadGain = mCalibObjects.mTPCPadGainCalib.get(); mustUpdate = true; } + + // NN clusterizer networks + if (mSpecConfig.nnLoadFromCCDB) { + + if (mSpecConfig.nnEvalMode[0] == "c1") { + pc.inputs().get("nn_classification_c1"); + } else if (mSpecConfig.nnEvalMode[0] == "c2") { + pc.inputs().get("nn_classification_c2"); + } + + pc.inputs().get("nn_regression_c1"); + if (mSpecConfig.nnEvalMode[1] == "r2") { + pc.inputs().get("nn_regression_c2"); + } + } } return mustUpdate; } diff --git a/macro/CMakeLists.txt b/macro/CMakeLists.txt index 843ad4a3be0ab..b5c51e50d3ffb 100644 --- a/macro/CMakeLists.txt +++ b/macro/CMakeLists.txt @@ -58,6 +58,7 @@ install(FILES CheckDigits_mft.C CreateGRPLHCIFObject.C getTimeStamp.C CreateSampleIRFrames.C + convert_onnx_to_root_serialized.C DESTINATION share/macro/) # FIXME: a lot of macros that are here should really be elsewhere. Those which @@ -149,6 +150,9 @@ o2_add_test_root_macro(checkTOFMatching.C O2::SimulationDataFormat O2::DataFormatsTOF) +o2_add_test_root_macro(convert_onnx_to_root_serialized.C + PUBLIC_LINK_LIBRARIES O2::GlobalTracking) + # FIXME: move to subsystem dir o2_add_test_root_macro(compareTopologyDistributions.C PUBLIC_LINK_LIBRARIES O2::DataFormatsITSMFT diff --git a/macro/convert_onnx_to_root_serialized.C b/macro/convert_onnx_to_root_serialized.C new file mode 100644 index 0000000000000..b1b8b981393a1 --- /dev/null +++ b/macro/convert_onnx_to_root_serialized.C @@ -0,0 +1,220 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file convert_onnx_to_root_serialized.C +/// \brief Utility functions to be executed as a ROOT macro for uploading ONNX models to CCDB as ROOT serialized objects and vice versa +/// \author Christian Sonnabend + +// Example execution: root -l -b -q '/scratch/csonnabe/MyO2/O2/GPU/GPUTracking/utils/convert_onnx_to_root_serialized.C("/scratch/csonnabe/PhD/jobs/clusterization/NN/output/21082025_smallWindow_clean/SC/training_data_21082025_reco_noise_supressed_p3t6_CoGselected/SC/PbPb_24arp2/0_5/class1/regression/399_noMom/network/net_fp16.onnx", "", 1, 1, "nnCCDBLayerType=FC/nnCCDBWithMomentum=0/inputDType=FP16/nnCCDBInteractionRate=500/outputDType=FP16/nnCCDBEvalType=regression_c1/nnCCDBBeamType=pp/partName=blob/quality=3", 1, 4108971600000, "Users/c/csonnabe/TPC/Clusterization", "model.root")' + +#include "ORTRootSerializer.h" +#include "CCDB/CcdbApi.h" +#include "CCDB/CcdbObjectInfo.h" +#include "TFile.h" +#include +#include + +o2::tpc::ORTRootSerializer serializer; + +/// Dumps the char* to a .onnx file -> Directly readable by ONNX runtime or Netron +void dumpOnnxToFile(const char* modelBuffer, uint32_t size, const std::string outputPath) +{ + std::ofstream outFile(outputPath, std::ios::binary | std::ios::trunc); + if (!outFile.is_open()) { + throw std::runtime_error("Failed to open output ONNX file: " + outputPath); + } + outFile.write(modelBuffer, static_cast(size)); + if (!outFile) { + throw std::runtime_error("Failed while writing data to: " + outputPath); + } + outFile.close(); +} + +/// Initialize the serialization from an ONNX file +void readOnnxModelFromFile(const std::string modelPath) +{ + std::ifstream inFile(modelPath, std::ios::binary | std::ios::ate); + if (!inFile.is_open()) { + throw std::runtime_error("Could not open input ONNX file " + modelPath); + } + std::streamsize size = inFile.tellg(); + std::vector mModelBuffer(size); + inFile.seekg(0, std::ios::beg); + if (!inFile.read(mModelBuffer.data(), size)) { + throw std::runtime_error("Could not read input ONNX file " + modelPath); + } + inFile.close(); + serializer.setOnnxModel(mModelBuffer.data(), static_cast(size)); +} + +/// Initialize the serialization from a ROOT file +void readRootModelFromFile(const std::string rootFilePath, std::string key) +{ + TFile inRootFile(rootFilePath.c_str()); + if (inRootFile.IsZombie()) { + throw std::runtime_error("Could not open input ROOT file " + rootFilePath); + } + auto* serPtr = inRootFile.Get(key.c_str()); + if (!serPtr) { + throw std::runtime_error("Could not find " + key + " in ROOT file " + rootFilePath); + } + serializer = *serPtr; + inRootFile.Close(); +} + +/// Serialize the ONNX model to a ROOT object and store to file +void onnxToRoot(std::string infile, std::string outfile, std::string key) +{ + readOnnxModelFromFile(infile); + TFile outRootFile(outfile.c_str(), "RECREATE"); + if (outRootFile.IsZombie()) { + throw std::runtime_error("Could not create output ROOT file " + outfile); + } + outRootFile.WriteObject(&serializer, key.c_str()); + outRootFile.Close(); +} + +/// Deserialize the ONNX model from a ROOT object and store to a .onnx file +void rootToOnnx(std::string infile, std::string outfile, std::string key) +{ + TFile inRootFile(infile.c_str()); + if (inRootFile.IsZombie()) { + throw std::runtime_error("Could not open input ROOT file " + infile); + } + auto* serPtr = inRootFile.Get(key.c_str()); + if (!serPtr) { + throw std::runtime_error("Could not find " + key + " in ROOT file " + infile); + } + serializer = *serPtr; + + std::ofstream outFile(outfile, std::ios::binary | std::ios::trunc); + if (!outFile.is_open()) { + throw std::runtime_error("Failed to open output ONNX file: " + outfile); + } + outFile.write(serializer.getONNXModel(), static_cast(serializer.getONNXModelSize())); + if (!outFile) { + throw std::runtime_error("Failed while writing data to: " + outfile); + } + outFile.close(); + + inRootFile.Close(); +} + +/// Upload the ONNX model to CCDB from an ONNX file +/// !!! Adjust the metadata, path and validity !!! +void uploadToCCDBFromONNX(std::string onnxFile, + const std::map& metadata, + // { // some example metadata entries + // "nnCCDBLayerType": "FC", + // "nnCCDBWithMomentum": "0", + // "inputDType": "FP16", + // "nnCCDBInteractionRate": "500", + // "outputDType": "FP16", + // "nnCCDBEvalType": "regression_c1", + // "nnCCDBBeamType": "pp", + // "partName": "blob", + // "quality": "3" + // } + long tsMin /* = 1 */, + long tsMax /* = 4108971600000 */, + std::string ccdbPath /* = "Users/c/csonnabe/TPC/Clusterization" */, + std::string objname /* = "net_regression_r1.root" */, + std::string ccdbUrl /* = "http://alice-ccdb.cern.ch" */) +{ + readOnnxModelFromFile(onnxFile); + + o2::ccdb::CcdbApi api; + api.init(ccdbUrl); + + // build full CCDB path including filename + const std::string fullPath = ccdbPath; //.back() == '/' ? (ccdbPath + objname) : (ccdbPath + "/" + objname); + + api.storeAsTFileAny(&serializer, fullPath, metadata, tsMin, tsMax); +} + +/// Upload the ONNX model to CCDB from a ROOT file +/// !!! Adjust the metadata, path and validity !!! +void uploadToCCDBFromROOT(std::string rootFile, + const std::map& metadata, + long tsMin /* = 1 */, + long tsMax /* = 4108971600000 */, + std::string ccdbPath /* = "Users/c/csonnabe/TPC/Clusterization" */, + std::string objname /* = "net_regression_r1.root" */, + std::string ccdbUrl /* = "http://alice-ccdb.cern.ch" */) +{ + // read ROOT file, extract ORTRootSerializer object and upload via storeAsTFileAny + TFile inRootFile(rootFile.c_str()); + if (inRootFile.IsZombie()) { + throw std::runtime_error("Could not open input ROOT file " + rootFile); + } + + // if objname is empty, fall back to default CCDB object key + const std::string key = objname.empty() ? o2::ccdb::CcdbApi::CCDBOBJECT_ENTRY : objname; + + auto* serPtr = inRootFile.Get(key.c_str()); + if (!serPtr) { + inRootFile.Close(); + throw std::runtime_error("Could not find " + key + " in ROOT file " + rootFile); + } + serializer = *serPtr; + + o2::ccdb::CcdbApi api; + api.init(ccdbUrl); + + // build full CCDB path including filename + const std::string fullPath = ccdbPath; //.back() == '/' ? (ccdbPath + objname) : (ccdbPath + "/" + objname); + + api.storeAsTFileAny(&serializer, fullPath, metadata, tsMin, tsMax); + + inRootFile.Close(); +} + +void convert_onnx_to_root_serialized(const std::string& onnxFile, + const std::string& rootFile, + int mode = 0, + int ccdbUpload = 0, + const std::string& metadataStr = "nnCCDBLayerType=FC/nnCCDBWithMomentum=0/inputDType=FP16/nnCCDBInteractionRate=500/outputDType=FP16/nnCCDBEvalType=regression_c1/nnCCDBBeamType=pp/partName=blob/quality=3", + long tsMin = 1, + long tsMax = 4108971600000, + std::string ccdbPath = "Users/c/csonnabe/TPC/Clusterization", + std::string objname = "net_regression_r1.root", + std::string ccdbUrl = "http://alice-ccdb.cern.ch") +{ + // parse metadataStr of the form key=value/key2=value2/... + std::map metadata; + std::size_t start = 0; + while (start < metadataStr.size()) { + auto sep = metadataStr.find('/', start); + auto token = metadataStr.substr(start, sep == std::string::npos ? std::string::npos : sep - start); + if (!token.empty()) { + auto eq = token.find('='); + if (eq != std::string::npos && eq > 0 && eq + 1 < token.size()) { + metadata.emplace(token.substr(0, eq), token.substr(eq + 1)); + } + } + if (sep == std::string::npos) { + break; + } + start = sep + 1; + } + + if (ccdbUpload == 0) { + if (mode == 0) + onnxToRoot(onnxFile, rootFile, o2::ccdb::CcdbApi::CCDBOBJECT_ENTRY); + else if (mode == 1) + rootToOnnx(rootFile, onnxFile, o2::ccdb::CcdbApi::CCDBOBJECT_ENTRY); + } else if (ccdbUpload == 1) { + if (mode == 0) + uploadToCCDBFromROOT(rootFile, metadata, tsMin, tsMax, ccdbPath, objname, ccdbUrl); + else if (mode == 1) + uploadToCCDBFromONNX(onnxFile, metadata, tsMin, tsMax, ccdbPath, objname, ccdbUrl); + } +}