Skip to content

Commit acbbe9d

Browse files
ktfktf
committed
Improve pidTPCML.cxx
* Out of line code to reduce memory usage while compiling pidTPC tasks * Change to use an array for the intermediate tracks representation. It would probably be better to read directly the values from the arrow tables, but that will come in a different PR.
1 parent e2f2256 commit acbbe9d

3 files changed

Lines changed: 232 additions & 197 deletions

File tree

‎Common/TableProducer/PID/CMakeLists.txt‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,11 @@
1111

1212
# TOF
1313

14+
15+
o2physics_add_library(PIDTPCCore
16+
SOURCES pidTPCML.cxx
17+
PUBLIC_LINK_LIBRARIES O2::Framework ONNXRuntime::ONNXRuntime
18+
)
1419
o2physics_add_dpl_workflow(pid-tof-base
1520
SOURCES pidTOFBase.cxx
1621
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2::TOFBase
@@ -35,12 +40,12 @@ o2physics_add_dpl_workflow(pid-tof-full
3540

3641
o2physics_add_dpl_workflow(pid-tpc
3742
SOURCES pidTPC.cxx
38-
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime
43+
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime
3944
COMPONENT_NAME Analysis)
4045

4146
o2physics_add_dpl_workflow(pid-tpc-full
4247
SOURCES pidTPCFull.cxx
43-
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore ONNXRuntime::ONNXRuntime
48+
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore O2Physics::PIDTPCCore ONNXRuntime::ONNXRuntime
4449
COMPONENT_NAME Analysis)
4550

4651
# HMPID
@@ -55,4 +60,4 @@ o2physics_add_dpl_workflow(pid-hmpid-qa
5560
o2physics_add_dpl_workflow(pid-bayes
5661
SOURCES pidBayes.cxx
5762
PUBLIC_LINK_LIBRARIES O2::Framework O2::DetectorsBase O2Physics::AnalysisCore
58-
COMPONENT_NAME Analysis)
63+
COMPONENT_NAME Analysis)

‎Common/TableProducer/PID/pidTPCML.cxx‎

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
///
13+
/// \file pidTPCML.h
14+
///
15+
/// \author Christian Sonnabend <christian.sonnabend@cern.ch>
16+
/// \author Nicolò Jacazio <nicolo.jacazio@cern.ch>
17+
///
18+
/// \brief A class for loading an ONNX neural network and evaluating it for the TPC PID response
19+
///
20+
#include <vector>
21+
22+
#include "TSystem.h"
23+
24+
// O2 includes
25+
#include "Framework/Logger.h"
26+
#include "Common/TableProducer/PID/pidTPCML.h"
27+
#include "ReconstructionDataFormats/PID.h"
28+
#include <onnxruntime/core/session/experimental_onnxruntime_cxx_api.h>
29+
30+
namespace o2::pid::tpc
31+
{
32+
33+
std::string Network::printShape(const std::vector<int64_t>& v)
34+
{
35+
std::stringstream ss("");
36+
for (size_t i = 0; i < v.size() - 1; i++)
37+
ss << v[i] << "x";
38+
ss << v[v.size() - 1];
39+
return ss.str();
40+
}
41+
42+
Network::Network(std::string path,
43+
bool enableOptimization = true)
44+
{
45+
46+
/*
47+
Constructor: Creating a class instance from a file and enabling optimizations with the boolean option.
48+
- Input:
49+
-- pathLocally: std::string ; Path to the model file;
50+
-- loadFromAlien: bool ; Download network from AliEn directory (true) or use local file (false)
51+
-- pathAlien: std::string ; if loadFromAlien is true, then the network will be downloaded from pathAlien
52+
-- enableOptimization: bool ; enabling optimizations for the loaded model in the session options;
53+
*/
54+
55+
LOG(info) << "--- Neural Network for the TPC PID response correction ---";
56+
57+
mEnv = std::make_shared<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "pid-neural-network");
58+
if (enableOptimization) {
59+
sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
60+
}
61+
62+
mSession.reset(new Ort::Experimental::Session{*mEnv, path, sessionOptions});
63+
64+
mInputNames = mSession->GetInputNames();
65+
mInputShapes = mSession->GetInputShapes();
66+
mOutputNames = mSession->GetOutputNames();
67+
mOutputShapes = mSession->GetOutputShapes();
68+
69+
LOG(info) << "Input Nodes:";
70+
for (size_t i = 0; i < mInputNames.size(); i++) {
71+
LOG(info) << "\t" << mInputNames[i] << " : " << printShape(mInputShapes[i]);
72+
}
73+
74+
LOG(info) << "Output Nodes:";
75+
for (size_t i = 0; i < mOutputNames.size(); i++) {
76+
LOG(info) << "\t" << mOutputNames[i] << " : " << printShape(mOutputShapes[i]);
77+
}
78+
79+
LOG(info) << "--- Network initialized! ---";
80+
81+
} // Network::Network(std::string, bool)
82+
83+
Network& Network::operator=(Network& inst)
84+
{
85+
86+
/*
87+
Operator: Setting instances of one class to the instances of an input class
88+
- Input:
89+
-- inst: const Network& ; An instance of the network class;
90+
- Output:
91+
-- *this: Network& ; An instance with the private properties of the "inst" (input) instance;
92+
*/
93+
94+
mEnv = inst.mEnv;
95+
mSession = inst.mSession;
96+
// sessionOptions = inst.sessionOptions; // Comment (Christian): Somehow the session options throw an error when trying to copy
97+
mInputNames = inst.mInputNames;
98+
mInputShapes = inst.mInputShapes;
99+
mOutputNames = inst.mOutputNames;
100+
mOutputShapes = inst.mOutputShapes;
101+
102+
LOG(debug) << "Network copied!";
103+
104+
return *this;
105+
106+
} // Network& Network::operator=(const Network &)
107+
108+
template <typename C, typename T>
109+
std::array<float, 6> Network::createInputFromTrack(const C& collision_it, const T& track, const uint8_t id) const
110+
{
111+
112+
/*
113+
Function: Creating a std::vector<float> from a track with the variables that the network has been trained on
114+
- Input:
115+
-- collisions_it const C& ; An iterator of a collisions table of the form: soa::Join<aod::Collisions, aod::Mults>::iterator const& collision
116+
-- track: const T& ; A track, typically from soa::Join<...> tables or of their iterators;
117+
-- id: uint8_t ; The id of a particle used for the mass assignment with o2::track::pid_constants::sMasses[id];
118+
- Output:
119+
-- inputValues: std::vector<float> ; A std::vector<float> with the input variables for the network;
120+
*/
121+
122+
const float p = track.tpcInnerParam();
123+
const float tgl = track.tgl();
124+
const float signed1Pt = track.signed1Pt();
125+
const float mass = o2::track::pid_constants::sMasses[id];
126+
const float multTPC = collision_it.multTPC() / 11000.;
127+
const float ncl = std::sqrt(159. / track.tpcNClsFound());
128+
129+
return {p, tgl, signed1Pt, mass, multTPC, ncl};
130+
}
131+
132+
std::vector<Ort::Value> Network::createTensor(std::array<float, 6> input) const
133+
{
134+
135+
/*
136+
Function: Creating a std::vector<Ort::Value> from a std::vector<float>
137+
- Input:
138+
-- input: std::vector<float> ; The vector from which the tensor should be created;
139+
- Output:
140+
-- inputValues: std::vector<Ort::Value> ; An ONNX tensor;
141+
*/
142+
143+
int64_t size = input.size();
144+
std::vector<int64_t> input_shape{size / mInputShapes[0][1], mInputShapes[0][1]};
145+
std::vector<Ort::Value> inputTensors;
146+
inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor<float>(input.data(), size, input_shape));
147+
148+
return inputTensors;
149+
150+
} // function std::vector<Ort::Value> Network::createTensor(std::vector<float>)
151+
152+
float* Network::evalNetwork(std::vector<Ort::Value> input)
153+
{
154+
155+
/*
156+
Function: Evaluating the network for a std::vector<Ort::Value>
157+
- Input:
158+
-- input: std::vector<Ort::Value> ; The tensor which should be evaluated by the network;
159+
- Output:
160+
-- output_values: const float* ; A float array which can be indexed;
161+
*/
162+
163+
try {
164+
LOG(debug) << "Shape of input (tensor): " << printShape(input[0].GetTensorTypeAndShapeInfo().GetShape());
165+
166+
auto outputTensors = mSession->Run(mInputNames, input, mOutputNames);
167+
float* output_values = outputTensors[0].GetTensorMutableData<float>();
168+
LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape());
169+
170+
// std::vector<float> output_vals(std::begin(output_values), std::end(output_values));
171+
172+
return output_values;
173+
174+
} catch (const Ort::Exception& exception) {
175+
LOG(error) << "Error running model inference: " << exception.what();
176+
177+
return 0;
178+
}
179+
180+
} // function Network::evalNetwork(std::vector<Ort::Value>)
181+
182+
float* Network::evalNetwork(std::vector<float> input)
183+
{
184+
185+
/*
186+
Function: Evaluating the network for a std::vector<float>
187+
- Input:
188+
-- input: std::vector<float> ; The vector which should be evaluated by the network.
189+
The network will evaluate n inputs where n = input.size()/input_nodes with input_nodes = #(input neurons);
190+
- Output:
191+
-- output_values: const float* ; A float array which can be indexed;
192+
*/
193+
194+
int64_t size = input.size();
195+
std::vector<int64_t> input_shape{size / mInputShapes[0][1], mInputShapes[0][1]};
196+
std::vector<Ort::Value> inputTensors;
197+
inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor<float>(input.data(), size, input_shape));
198+
199+
try {
200+
201+
LOG(debug) << "Shape of input (vector): " << printShape(input_shape);
202+
auto outputTensors = mSession->Run(mInputNames, inputTensors, mOutputNames);
203+
LOG(debug) << "Shape of output (tensor): " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape());
204+
float* output_values = outputTensors[0].GetTensorMutableData<float>();
205+
206+
return output_values;
207+
208+
} catch (const Ort::Exception& exception) {
209+
210+
LOG(error) << "Error running model inference: " << exception.what();
211+
212+
return 0;
213+
}
214+
215+
} // function Network::evalNetwork(std::vector<float>)
216+
217+
} // namespace o2::pid::tpc

0 commit comments

Comments
 (0)