diff --git a/Common/TableProducer/PID/CMakeLists.txt b/Common/TableProducer/PID/CMakeLists.txt index 82bad53ddaf..f8704c9bfab 100644 --- a/Common/TableProducer/PID/CMakeLists.txt +++ b/Common/TableProducer/PID/CMakeLists.txt @@ -11,6 +11,11 @@ # TOF + +o2physics_add_library(PIDTPCCore + SOURCES pidTPCML.cxx + PUBLIC_LINK_LIBRARIES O2::Framework ONNXRuntime::ONNXRuntime +) o2physics_add_dpl_workflow(pid-tof-base SOURCES pidTOFBase.cxx PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2::TOFBase @@ -35,12 +40,12 @@ o2physics_add_dpl_workflow(pid-tof-full o2physics_add_dpl_workflow(pid-tpc SOURCES pidTPC.cxx - PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime + PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(pid-tpc-full SOURCES pidTPCFull.cxx - PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime + PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime COMPONENT_NAME Analysis) # HMPID @@ -55,4 +60,4 @@ o2physics_add_dpl_workflow(pid-hmpid-qa o2physics_add_dpl_workflow(pid-bayes SOURCES pidBayes.cxx PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore - COMPONENT_NAME Analysis) \ No newline at end of file + COMPONENT_NAME Analysis) diff --git a/Common/TableProducer/PID/pidTPCML.cxx b/Common/TableProducer/PID/pidTPCML.cxx new file mode 100644 index 00000000000..6207065f1ec --- /dev/null +++ b/Common/TableProducer/PID/pidTPCML.cxx @@ -0,0 +1,217 @@ +// 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 pidTPCML.h +/// +/// \author Christian Sonnabend +/// \author Nicolò Jacazio +/// +/// \brief A class for loading an ONNX neural network and evaluating it for the TPC PID response +/// +#include + +#include "TSystem.h" + +// O2 includes +#include "Framework/Logger.h" +#include "Common/TableProducer/PID/pidTPCML.h" +#include "ReconstructionDataFormats/PID.h" +#include + +namespace o2::pid::tpc +{ + +std::string Network::printShape(const std::vector& v) +{ + std::stringstream ss(""); + for (size_t i = 0; i < v.size() - 1; i++) + ss << v[i] << "x"; + ss << v[v.size() - 1]; + return ss.str(); +} + +Network::Network(std::string path, + bool enableOptimization = true) +{ + + /* + Constructor: Creating a class instance from a file and enabling optimizations with the boolean option. + - Input: + -- pathLocally: std::string ; Path to the model file; + -- loadFromAlien: bool ; Download network from AliEn directory (true) or use local file (false) + -- pathAlien: std::string ; if loadFromAlien is true, then the network will be downloaded from pathAlien + -- enableOptimization: bool ; enabling optimizations for the loaded model in the session options; + */ + + LOG(info) << "--- Neural Network for the TPC PID response correction ---"; + + mEnv = std::make_shared(ORT_LOGGING_LEVEL_WARNING, "pid-neural-network"); + if (enableOptimization) { + sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); + } + + mSession.reset(new Ort::Experimental::Session{*mEnv, path, sessionOptions}); + + mInputNames = mSession->GetInputNames(); + mInputShapes = mSession->GetInputShapes(); + mOutputNames = mSession->GetOutputNames(); + mOutputShapes = mSession->GetOutputShapes(); + + LOG(info) << "Input Nodes:"; + for (size_t i = 0; i < mInputNames.size(); i++) { + LOG(info) << "\t" << mInputNames[i] << " : " << printShape(mInputShapes[i]); + } + + LOG(info) << "Output Nodes:"; + for (size_t i = 0; i < mOutputNames.size(); i++) { + LOG(info) << "\t" << mOutputNames[i] << " : " << printShape(mOutputShapes[i]); + } + + LOG(info) << "--- Network initialized! ---"; + +} // Network::Network(std::string, bool) + +Network& Network::operator=(Network& inst) +{ + + /* + Operator: Setting instances of one class to the instances of an input class + - Input: + -- inst: const Network& ; An instance of the network class; + - Output: + -- *this: Network& ; An instance with the private properties of the "inst" (input) instance; + */ + + mEnv = inst.mEnv; + mSession = inst.mSession; + // sessionOptions = inst.sessionOptions; // Comment (Christian): Somehow the session options throw an error when trying to copy + mInputNames = inst.mInputNames; + mInputShapes = inst.mInputShapes; + mOutputNames = inst.mOutputNames; + mOutputShapes = inst.mOutputShapes; + + LOG(debug) << "Network copied!"; + + return *this; + +} // Network& Network::operator=(const Network &) + +template +std::array Network::createInputFromTrack(const C& collision_it, const T& track, const uint8_t id) const +{ + + /* + Function: Creating a std::vector from a track with the variables that the network has been trained on + - Input: + -- collisions_it const C& ; An iterator of a collisions table of the form: soa::Join::iterator const& collision + -- track: const T& ; A track, typically from soa::Join<...> tables or of their iterators; + -- id: uint8_t ; The id of a particle used for the mass assignment with o2::track::pid_constants::sMasses[id]; + - Output: + -- inputValues: std::vector ; A std::vector with the input variables for the network; + */ + + const float p = track.tpcInnerParam(); + const float tgl = track.tgl(); + const float signed1Pt = track.signed1Pt(); + const float mass = o2::track::pid_constants::sMasses[id]; + const float multTPC = collision_it.multTPC() / 11000.; + const float ncl = std::sqrt(159. / track.tpcNClsFound()); + + return {p, tgl, signed1Pt, mass, multTPC, ncl}; +} + +std::vector Network::createTensor(std::array input) const +{ + + /* + Function: Creating a std::vector from a std::vector + - Input: + -- input: std::vector ; The vector from which the tensor should be created; + - Output: + -- inputValues: std::vector ; An ONNX tensor; + */ + + int64_t size = input.size(); + std::vector input_shape{size / mInputShapes[0][1], mInputShapes[0][1]}; + std::vector inputTensors; + inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor(input.data(), size, input_shape)); + + return inputTensors; + +} // function std::vector Network::createTensor(std::vector) + +float* Network::evalNetwork(std::vector input) +{ + + /* + Function: Evaluating the network for a std::vector + - Input: + -- input: std::vector ; The tensor which should be evaluated by the network; + - Output: + -- output_values: const float* ; A float array which can be indexed; + */ + + try { + LOG(debug) << "Shape of input (tensor): " << printShape(input[0].GetTensorTypeAndShapeInfo().GetShape()); + + auto outputTensors = mSession->Run(mInputNames, input, mOutputNames); + float* output_values = outputTensors[0].GetTensorMutableData(); + LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape()); + + // std::vector output_vals(std::begin(output_values), std::end(output_values)); + + return output_values; + + } catch (const Ort::Exception& exception) { + LOG(error) << "Error running model inference: " << exception.what(); + + return 0; + } + +} // function Network::evalNetwork(std::vector) + +float* Network::evalNetwork(std::vector input) +{ + + /* + Function: Evaluating the network for a std::vector + - Input: + -- input: std::vector ; The vector which should be evaluated by the network. + The network will evaluate n inputs where n = input.size()/input_nodes with input_nodes = #(input neurons); + - Output: + -- output_values: const float* ; A float array which can be indexed; + */ + + int64_t size = input.size(); + std::vector input_shape{size / mInputShapes[0][1], mInputShapes[0][1]}; + std::vector inputTensors; + inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor(input.data(), size, input_shape)); + + try { + + LOG(debug) << "Shape of input (vector): " << printShape(input_shape); + auto outputTensors = mSession->Run(mInputNames, inputTensors, mOutputNames); + LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape()); + float* output_values = outputTensors[0].GetTensorMutableData(); + + return output_values; + + } catch (const Ort::Exception& exception) { + + LOG(error) << "Error running model inference: " << exception.what(); + + return 0; + } + +} // function Network::evalNetwork(std::vector) + +} // namespace o2::pid::tpc diff --git a/Common/TableProducer/PID/pidTPCML.h b/Common/TableProducer/PID/pidTPCML.h index 9d4bd924f4c..245989b80a7 100644 --- a/Common/TableProducer/PID/pidTPCML.h +++ b/Common/TableProducer/PID/pidTPCML.h @@ -21,14 +21,12 @@ #ifndef O2_PID_TPC_ML_H_ #define O2_PID_TPC_ML_H_ -#include - -#include "TSystem.h" - // O2 includes -#include "Framework/Logger.h" #include "ReconstructionDataFormats/PID.h" #include +#include +#include +#include namespace o2::pid::tpc { @@ -46,8 +44,8 @@ class Network // Functions template - std::vector createInputFromTrack(const C&, const T&, const uint8_t) const; // create a std::vector with all the inputs for the network - std::vector createTensor(std::vector) const; // create a std::vector (= ONNX tensor) for model input + std::array createInputFromTrack(const C&, const T&, const uint8_t) const; // create a std::vector with all the inputs for the network + std::vector createTensor(std::array) const; // create a std::vector (= ONNX tensor) for model input float* evalNetwork(std::vector); // evaluate the network on a std::vector (= ONNX tensor) float* evalNetwork(std::vector); // evaluate the network on a std::vector @@ -68,198 +66,13 @@ class Network std::vector> mOutputShapes; // Internal function for printing the shape of tensors: See https://github.com/saganatt/PID_ML_in_O2 or O2Physics/Tools/PIDML/simpleApplyPidOnnxModel.cxx - std::string printShape(const std::vector& v) - { - std::stringstream ss(""); - for (size_t i = 0; i < v.size() - 1; i++) - ss << v[i] << "x"; - ss << v[v.size() - 1]; - return ss.str(); - } + std::string printShape(const std::vector& v); // Class version ClassDefNV(Network, 2); }; // class Network -Network::Network(std::string path, - bool enableOptimization = true) -{ - - /* - Constructor: Creating a class instance from a file and enabling optimizations with the boolean option. - - Input: - -- pathLocally: std::string ; Path to the model file; - -- loadFromAlien: bool ; Download network from AliEn directory (true) or use local file (false) - -- pathAlien: std::string ; if loadFromAlien is true, then the network will be downloaded from pathAlien - -- enableOptimization: bool ; enabling optimizations for the loaded model in the session options; - */ - - LOG(info) << "--- Neural Network for the TPC PID response correction ---"; - - mEnv = std::make_shared(ORT_LOGGING_LEVEL_WARNING, "pid-neural-network"); - if (enableOptimization) { - sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED); - } - - mSession.reset(new Ort::Experimental::Session{*mEnv, path, sessionOptions}); - - mInputNames = mSession->GetInputNames(); - mInputShapes = mSession->GetInputShapes(); - mOutputNames = mSession->GetOutputNames(); - mOutputShapes = mSession->GetOutputShapes(); - - LOG(info) << "Input Nodes:"; - for (size_t i = 0; i < mInputNames.size(); i++) { - LOG(info) << "\t" << mInputNames[i] << " : " << printShape(mInputShapes[i]); - } - - LOG(info) << "Output Nodes:"; - for (size_t i = 0; i < mOutputNames.size(); i++) { - LOG(info) << "\t" << mOutputNames[i] << " : " << printShape(mOutputShapes[i]); - } - - LOG(info) << "--- Network initialized! ---"; - -} // Network::Network(std::string, bool) - -Network& Network::operator=(Network& inst) -{ - - /* - Operator: Setting instances of one class to the instances of an input class - - Input: - -- inst: const Network& ; An instance of the network class; - - Output: - -- *this: Network& ; An instance with the private properties of the "inst" (input) instance; - */ - - mEnv = inst.mEnv; - mSession = inst.mSession; - // sessionOptions = inst.sessionOptions; // Comment (Christian): Somehow the session options throw an error when trying to copy - mInputNames = inst.mInputNames; - mInputShapes = inst.mInputShapes; - mOutputNames = inst.mOutputNames; - mOutputShapes = inst.mOutputShapes; - - LOG(debug) << "Network copied!"; - - return *this; - -} // Network& Network::operator=(const Network &) - -template -std::vector Network::createInputFromTrack(const C& collision_it, const T& track, const uint8_t id) const -{ - - /* - Function: Creating a std::vector from a track with the variables that the network has been trained on - - Input: - -- collisions_it const C& ; An iterator of a collisions table of the form: soa::Join::iterator const& collision - -- track: const T& ; A track, typically from soa::Join<...> tables or of their iterators; - -- id: uint8_t ; The id of a particle used for the mass assignment with o2::track::pid_constants::sMasses[id]; - - Output: - -- inputValues: std::vector ; A std::vector with the input variables for the network; - */ - - const float p = track.tpcInnerParam(); - const float tgl = track.tgl(); - const float signed1Pt = track.signed1Pt(); - const float mass = o2::track::pid_constants::sMasses[id]; - const float multTPC = collision_it.multTPC() / 11000.; - const float ncl = std::sqrt(159. / track.tpcNClsFound()); - - std::vector inputValues{p, tgl, signed1Pt, mass, multTPC, ncl}; - - return inputValues; - -} // std::vector Network::createInputFromTrack(const C&, const T&, uint8_t) - -std::vector Network::createTensor(std::vector input) const -{ - - /* - Function: Creating a std::vector from a std::vector - - Input: - -- input: std::vector ; The vector from which the tensor should be created; - - Output: - -- inputValues: std::vector ; An ONNX tensor; - */ - - int64_t size = input.size(); - std::vector input_shape{size / mInputShapes[0][1], mInputShapes[0][1]}; - std::vector inputTensors; - inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor(input.data(), size, input_shape)); - - return inputTensors; - -} // function std::vector Network::createTensor(std::vector) - -float* Network::evalNetwork(std::vector input) -{ - - /* - Function: Evaluating the network for a std::vector - - Input: - -- input: std::vector ; The tensor which should be evaluated by the network; - - Output: - -- output_values: const float* ; A float array which can be indexed; - */ - - try { - LOG(debug) << "Shape of input (tensor): " << printShape(input[0].GetTensorTypeAndShapeInfo().GetShape()); - - auto outputTensors = mSession->Run(mInputNames, input, mOutputNames); - float* output_values = outputTensors[0].GetTensorMutableData(); - LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape()); - - // std::vector output_vals(std::begin(output_values), std::end(output_values)); - - return output_values; - - } catch (const Ort::Exception& exception) { - LOG(error) << "Error running model inference: " << exception.what(); - - return 0; - } - -} // function Network::evalNetwork(std::vector) - -float* Network::evalNetwork(std::vector input) -{ - - /* - Function: Evaluating the network for a std::vector - - Input: - -- input: std::vector ; The vector which should be evaluated by the network. - The network will evaluate n inputs where n = input.size()/input_nodes with input_nodes = #(input neurons); - - Output: - -- output_values: const float* ; A float array which can be indexed; - */ - - int64_t size = input.size(); - std::vector input_shape{size / mInputShapes[0][1], mInputShapes[0][1]}; - std::vector inputTensors; - inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor(input.data(), size, input_shape)); - - try { - - LOG(debug) << "Shape of input (vector): " << printShape(input_shape); - auto outputTensors = mSession->Run(mInputNames, inputTensors, mOutputNames); - LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape()); - float* output_values = outputTensors[0].GetTensorMutableData(); - - return output_values; - - } catch (const Ort::Exception& exception) { - - LOG(error) << "Error running model inference: " << exception.what(); - - return 0; - } - -} // function Network::evalNetwork(std::vector) - } // namespace o2::pid::tpc -#endif // O2_PID_TPC_ML_H_ \ No newline at end of file +#endif // O2_PID_TPC_ML_H_