Skip to content

Commit 5fef658

Browse files
authored
Merge pull request #48249 from SegmentLinking/t5_embed_new
Track Embeddings for Improved Duplicate Removal in LST
2 parents f79df57 + ab51b65 commit 5fef658

File tree

16 files changed

+2761
-53
lines changed

16 files changed

+2761
-53
lines changed

RecoTracker/LSTCore/interface/Common.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,9 @@ namespace lst {
6464
};
6565
struct Params_pLS {
6666
static constexpr int kLayers = 2, kHits = 4;
67+
static constexpr int kEmbed = 6;
6768
using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
69+
using ArrayFxEmbed = edm::StdArray<float, kEmbed>;
6870
};
6971
struct Params_LS {
7072
static constexpr int kLayers = 2, kHits = 4;
@@ -84,9 +86,11 @@ namespace lst {
8486
};
8587
struct Params_T5 {
8688
static constexpr int kLayers = 5, kHits = 10;
89+
static constexpr int kEmbed = 6;
8790
using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
8891
using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
8992
using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
93+
using ArrayFxEmbed = edm::StdArray<float, kEmbed>;
9094
};
9195
struct Params_pT5 {
9296
static constexpr int kLayers = 7, kHits = 14;

RecoTracker/LSTCore/interface/PixelSegmentsSoA.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ namespace lst {
1010

1111
GENERATE_SOA_LAYOUT(PixelSegmentsSoALayout,
1212
SOA_COLUMN(uint4, pLSHitsIdxs),
13+
SOA_COLUMN(Params_pLS::ArrayFxEmbed, plsEmbed),
1314
SOA_COLUMN(char, isDup),
1415
SOA_COLUMN(bool, partOfPT5),
1516
SOA_COLUMN(float, score),

RecoTracker/LSTCore/interface/QuintupletsSoA.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ namespace lst {
1313
SOA_COLUMN(Params_T5::ArrayU16xLayers, lowerModuleIndices), // lower module index in each layer
1414
SOA_COLUMN(Params_T5::ArrayU8xLayers, logicalLayers), // layer ID
1515
SOA_COLUMN(Params_T5::ArrayUxHits, hitIndices), // hit indices
16+
SOA_COLUMN(Params_T5::ArrayFxEmbed, t5Embed), // t5 embedding vector
1617
SOA_COLUMN(FPX, innerRadius), // inner triplet circle radius
1718
SOA_COLUMN(FPX, bridgeRadius), // "middle"/bridge triplet radius
1819
SOA_COLUMN(FPX, outerRadius), // outer triplet radius

RecoTracker/LSTCore/interface/TripletsSoA.h

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,13 @@ namespace lst {
1313
SOA_COLUMN(Params_T3::ArrayU16xLayers, lowerModuleIndices), // lower module index in each layer
1414
SOA_COLUMN(Params_T3::ArrayU8xLayers, logicalLayers), // layer ID
1515
SOA_COLUMN(Params_T3::ArrayUxHits, hitIndices), // hit indices
16-
SOA_COLUMN(FPX, betaIn), // beta/chord angle of the inner segment
17-
SOA_COLUMN(float, centerX), // lower/anchor-hit based circle center x
18-
SOA_COLUMN(float, centerY), // lower/anchor-hit based circle center y
19-
SOA_COLUMN(float, radius), // lower/anchor-hit based circle radius
16+
SOA_COLUMN(FPX, betaIn), // beta/chord angle of the inner segment
17+
SOA_COLUMN(float, centerX), // lower/anchor-hit based circle center x
18+
SOA_COLUMN(float, centerY), // lower/anchor-hit based circle center y
19+
SOA_COLUMN(float, radius), // lower/anchor-hit based circle radius
20+
SOA_COLUMN(float, fakeScore), // DNN confidence score for fake t3
21+
SOA_COLUMN(float, promptScore), // DNN confidence score for real (prompt) t3
22+
SOA_COLUMN(float, displacedScore), // DNN confidence score for real (displaced) t3
2023
#ifdef CUT_VALUE_DEBUG
2124
SOA_COLUMN(float, zOut),
2225
SOA_COLUMN(float, rtOut),

RecoTracker/LSTCore/interface/alpaka/Common.h

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -58,34 +58,44 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
5858
namespace dnn {
5959

6060
// Common constants for both DNNs
61-
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
6261
HOST_DEVICE_CONSTANT float kPhi_norm = kPi;
6362
HOST_DEVICE_CONSTANT float kEtaSize = 0.25f; // Bin size in eta.
6463
constexpr unsigned int kPtBins = 2;
6564
constexpr unsigned int kEtaBins = 10;
6665

66+
namespace plsembdnn {
67+
HOST_DEVICE_CONSTANT float kEta_norm = 4.0f;
68+
HOST_DEVICE_CONSTANT float kEtaErr_norm = 0.00139f;
69+
HOST_DEVICE_CONSTANT float kWP[kEtaBins] = {
70+
0.9235f, 0.8974f, 0.9061f, 0.9431f, 0.8262f, 0.7998f, 0.7714f, 0.7017f, 0.6749f, 0.6624f};
71+
} // namespace plsembdnn
72+
6773
namespace t3dnn {
74+
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
6875
HOST_DEVICE_CONSTANT float kZ_max = 224.149505f;
6976
HOST_DEVICE_CONSTANT float kR_max = 98.932365f;
77+
HOST_DEVICE_CONSTANT unsigned int kOutputFeatures = 3;
7078
HOST_DEVICE_CONSTANT float kWp_prompt[kPtBins][kEtaBins] = {
71-
{0.4957, 0.5052, 0.5201, 0.5340, 0.4275, 0.4708, 0.4890, 0.4932, 0.5400, 0.5449},
72-
{0.0302, 0.0415, 0.0994, 0.1791, 0.1960, 0.2467, 0.3227, 0.3242, 0.2367, 0.2187}};
79+
{0.4957f, 0.5052f, 0.5201f, 0.5340f, 0.4275f, 0.4708f, 0.4890f, 0.4932f, 0.5400f, 0.5449f},
80+
{0.0302f, 0.0415f, 0.0994f, 0.1791f, 0.1960f, 0.2467f, 0.3227f, 0.3242f, 0.2367f, 0.2187f}};
7381
HOST_DEVICE_CONSTANT float kWp_displaced[kPtBins][kEtaBins] = {
74-
{0.0334, 0.0504, 0.0748, 0.0994, 0.1128, 0.1123, 0.1118, 0.1525, 0.1867, 0.1847},
75-
{0.0091, 0.0075, 0.0350, 0.0213, 0.0435, 0.0676, 0.1957, 0.1649, 0.1080, 0.1046}};
82+
{0.0334f, 0.0504f, 0.0748f, 0.0994f, 0.1128f, 0.1123f, 0.1118f, 0.1525f, 0.1867f, 0.1847f},
83+
{0.0091f, 0.0075f, 0.0350f, 0.0213f, 0.0435f, 0.0676f, 0.1957f, 0.1649f, 0.1080f, 0.1046f}};
7684
} // namespace t3dnn
7785

7886
namespace t5dnn {
87+
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
7988
HOST_DEVICE_CONSTANT float kZ_max = 267.2349854f;
8089
HOST_DEVICE_CONSTANT float kR_max = 110.1099396f;
8190
HOST_DEVICE_CONSTANT float kWp[kPtBins][kEtaBins] = {
82-
{0.4493, 0.4939, 0.5715, 0.6488, 0.5709, 0.5938, 0.7164, 0.7565, 0.8103, 0.8593},
83-
{0.4488, 0.4448, 0.5067, 0.5929, 0.4836, 0.4112, 0.4968, 0.4403, 0.5597, 0.5067}};
91+
{0.4493f, 0.4939f, 0.5715f, 0.6488f, 0.5709f, 0.5938f, 0.7164f, 0.7565f, 0.8103f, 0.8593f},
92+
{0.4488f, 0.4448f, 0.5067f, 0.5929f, 0.4836f, 0.4112f, 0.4968f, 0.4403f, 0.5597f, 0.5067f}};
8493
} // namespace t5dnn
8594

8695
namespace pt3dnn {
96+
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
8797
HOST_DEVICE_CONSTANT float kWp[kEtaBins] = {
88-
0.189, 0.1805, 0.2267, 0.3104, 0.4719, 0.3159, 0.1372, 0.1571, 0.3198, 0.186};
98+
0.189f, 0.1805f, 0.2267f, 0.3104f, 0.4719f, 0.3159f, 0.1372f, 0.1571f, 0.3198f, 0.186f};
8999
HOST_DEVICE_CONSTANT float kWpHigh = 0.0473f;
90100
} // namespace pt3dnn
91101

RecoTracker/LSTCore/src/alpaka/Kernels.h

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
#define RecoTracker_LSTCore_src_alpaka_Kernels_h
33

44
#include "HeterogeneousCore/AlpakaInterface/interface/workdivision.h"
5+
#include "FWCore/Utilities/interface/CMSUnrollLoop.h"
56

67
#include "RecoTracker/LSTCore/interface/alpaka/Common.h"
78
#include "RecoTracker/LSTCore/interface/ModulesSoA.h"
@@ -249,7 +250,15 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
249250
float dR2 = dEta * dEta + dPhi * dPhi;
250251
int nMatched = checkHitsT5(ix, jx, quintuplets);
251252
const int minNHitsForDup_T5 = 5;
252-
if (dR2 < 0.001f || nMatched >= minNHitsForDup_T5) {
253+
254+
float d2 = 0.f;
255+
CMS_UNROLL_LOOP
256+
for (unsigned int k = 0; k < Params_T5::kEmbed; ++k) {
257+
float diff = quintuplets.t5Embed()[ix][k] - quintuplets.t5Embed()[jx][k];
258+
d2 += diff * diff;
259+
}
260+
261+
if (((dR2 < 0.001f || nMatched >= minNHitsForDup_T5) && d2 < 1.0f) || (dR2 < 0.02f && d2 < 0.1f)) {
253262
if (isPT5_jx || score_rphisum1 > score_rphisum2) {
254263
rmQuintupletFromMemory(quintuplets, ix, true);
255264
} else if (isPT5_ix || score_rphisum1 < score_rphisum2) {

RecoTracker/LSTCore/src/alpaka/NeuralNetwork.h

Lines changed: 151 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
#include "T5NeuralNetworkWeights.h"
1111
#include "T3NeuralNetworkWeights.h"
1212
#include "pT3NeuralNetworkWeights.h"
13+
#include "T5EmbedNetworkWeights.h"
14+
#include "pLSEmbedNetworkWeights.h"
1315

1416
namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
1517

@@ -66,11 +68,11 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
6668
const unsigned int mdIndex2,
6769
const unsigned int mdIndex3,
6870
const float radius,
69-
const float betaIn) {
71+
const float betaIn,
72+
float (&output)[dnn::t3dnn::kOutputFeatures]) {
7073
// Constants for T3 DNN
7174
constexpr unsigned int kInputFeatures = 14;
7275
constexpr unsigned int kHiddenFeatures = 32;
73-
constexpr unsigned int kOutputFeatures = 3;
7476

7577
// Extract hit information
7678
float eta1 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex1]); // inner T3 anchor hit 1 eta
@@ -91,7 +93,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
9193

9294
// Build input feature vector matching training order
9395
float x[kInputFeatures] = {
94-
eta1 / dnn::kEta_norm, // First hit eta normalized
96+
eta1 / dnn::t3dnn::kEta_norm, // First hit eta normalized
9597
alpaka::math::abs(acc, phi1) / dnn::kPhi_norm, // First hit phi normalized
9698
z1 / dnn::t3dnn::kZ_max, // First hit z normalized
9799
r1 / dnn::t3dnn::kR_max, // First hit r normalized
@@ -112,7 +114,6 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
112114

113115
float x_1[kHiddenFeatures]; // Layer 1 output
114116
float x_2[kHiddenFeatures]; // Layer 2 output
115-
float x_3[kOutputFeatures]; // Layer 3 output (3 classes)
116117

117118
// Layer 1: Linear + Relu
118119
linear_layer<kInputFeatures, kHiddenFeatures>(x, x_1, dnn::t3dnn::wgtT_layer1, dnn::t3dnn::bias_layer1);
@@ -123,17 +124,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
123124
relu_activation<kHiddenFeatures>(x_2);
124125

125126
// Layer 3: Linear + Softmax
126-
linear_layer<kHiddenFeatures, kOutputFeatures>(
127-
x_2, x_3, dnn::t3dnn::wgtT_output_layer, dnn::t3dnn::bias_output_layer);
128-
softmax_activation<kOutputFeatures>(acc, x_3);
127+
linear_layer<kHiddenFeatures, dnn::t3dnn::kOutputFeatures>(
128+
x_2, output, dnn::t3dnn::wgtT_output_layer, dnn::t3dnn::bias_output_layer);
129+
softmax_activation<dnn::t3dnn::kOutputFeatures>(acc, output);
129130

130131
// Get pt and eta bin indices
131132
float t3_pt = radius * lst::k2Rinv1GeVf * 2;
132133
uint8_t pt_index = (t3_pt > 5);
133134
uint8_t bin_index = (eta1 > 2.5f) ? (dnn::kEtaBins - 1) : static_cast<unsigned int>(eta1 / dnn::kEtaSize);
134135

135-
return x_3[1] > dnn::t3dnn::kWp_prompt[pt_index][bin_index] ||
136-
x_3[2] > dnn::t3dnn::kWp_displaced[pt_index][bin_index];
136+
return output[1] > dnn::t3dnn::kWp_prompt[pt_index][bin_index] ||
137+
output[2] > dnn::t3dnn::kWp_displaced[pt_index][bin_index];
137138
}
138139
} // namespace t3dnn
139140

@@ -157,7 +158,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
157158
alpaka::math::log10(acc, pixelRadius),
158159
alpaka::math::log10(acc, pixRadiusError),
159160
alpaka::math::log10(acc, rzChiSquared),
160-
alpaka::math::abs(acc, pixelEta) / dnn::kEta_norm};
161+
alpaka::math::abs(acc, pixelEta) / dnn::pt3dnn::kEta_norm};
161162

162163
float x1[kHiddenFeatures];
163164
float x2[kHiddenFeatures];
@@ -227,7 +228,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
227228

228229
// Build the input feature vector using pairwise differences after the first hit
229230
float x[kInputFeatures] = {
230-
eta1 / dnn::kEta_norm, // inner T3: First hit eta normalized
231+
eta1 / dnn::t5dnn::kEta_norm, // inner T3: First hit eta normalized
231232
alpaka::math::abs(acc, phi1) / dnn::kPhi_norm, // inner T3: First hit phi normalized
232233
z1 / dnn::t5dnn::kZ_max, // inner T3: First hit z normalized
233234
r1 / dnn::t5dnn::kR_max, // inner T3: First hit r normalized
@@ -288,6 +289,145 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
288289
}
289290
} // namespace t5dnn
290291

292+
namespace t5embdnn {
293+
template <typename TAcc>
294+
ALPAKA_FN_ACC ALPAKA_FN_INLINE void runEmbed(TAcc const& acc,
295+
MiniDoubletsConst mds,
296+
unsigned int mdIndex1,
297+
unsigned int mdIndex2,
298+
unsigned int mdIndex3,
299+
unsigned int mdIndex4,
300+
unsigned int mdIndex5,
301+
float innerRadius,
302+
float outerRadius,
303+
float bridgeRadius,
304+
float fakeScore1,
305+
float promptScore1,
306+
float dispScore1,
307+
float fakeScore2,
308+
float promptScore2,
309+
float dispScore2,
310+
float (&embedding)[Params_T5::kEmbed]) {
311+
constexpr unsigned int kInputFeatures = 30;
312+
constexpr unsigned int kHiddenFeatures = 32;
313+
314+
float eta1 = mds.anchorEta()[mdIndex1];
315+
float eta2 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex2]);
316+
float eta3 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex3]);
317+
float eta4 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex4]);
318+
float eta5 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex5]);
319+
320+
float phi1 = mds.anchorPhi()[mdIndex1];
321+
float phi2 = mds.anchorPhi()[mdIndex2];
322+
float phi3 = mds.anchorPhi()[mdIndex3];
323+
float phi4 = mds.anchorPhi()[mdIndex4];
324+
float phi5 = mds.anchorPhi()[mdIndex5];
325+
326+
float z1 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex1]);
327+
float z2 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex2]);
328+
float z3 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex3]);
329+
float z4 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex4]);
330+
float z5 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex5]);
331+
332+
float r1 = mds.anchorRt()[mdIndex1];
333+
float r2 = mds.anchorRt()[mdIndex2];
334+
float r3 = mds.anchorRt()[mdIndex3];
335+
float r4 = mds.anchorRt()[mdIndex4];
336+
float r5 = mds.anchorRt()[mdIndex5];
337+
338+
float x[kInputFeatures] = {eta1 / dnn::t5dnn::kEta_norm,
339+
alpaka::math::cos(acc, phi1),
340+
alpaka::math::sin(acc, phi1),
341+
z1 / dnn::t5dnn::kZ_max,
342+
r1 / dnn::t5dnn::kR_max,
343+
344+
eta2 - alpaka::math::abs(acc, eta1),
345+
cms::alpakatools::deltaPhi(acc, phi2, phi1),
346+
(z2 - z1) / dnn::t5dnn::kZ_max,
347+
(r2 - r1) / dnn::t5dnn::kR_max,
348+
349+
eta3 - eta2,
350+
cms::alpakatools::deltaPhi(acc, phi3, phi2),
351+
(z3 - z2) / dnn::t5dnn::kZ_max,
352+
(r3 - r2) / dnn::t5dnn::kR_max,
353+
354+
eta4 - eta3,
355+
cms::alpakatools::deltaPhi(acc, phi4, phi3),
356+
(z4 - z3) / dnn::t5dnn::kZ_max,
357+
(r4 - r3) / dnn::t5dnn::kR_max,
358+
359+
eta5 - eta4,
360+
cms::alpakatools::deltaPhi(acc, phi5, phi4),
361+
(z5 - z4) / dnn::t5dnn::kZ_max,
362+
(r5 - r4) / dnn::t5dnn::kR_max,
363+
364+
1.0f / innerRadius,
365+
1.0f / bridgeRadius,
366+
1.0f / outerRadius,
367+
368+
fakeScore1,
369+
promptScore1,
370+
dispScore1,
371+
(fakeScore2 - fakeScore1),
372+
(promptScore2 - promptScore1),
373+
(dispScore2 - dispScore1)};
374+
375+
float h1[kHiddenFeatures];
376+
float h2[kHiddenFeatures];
377+
378+
linear_layer<kInputFeatures, kHiddenFeatures>(x, h1, dnn::t5embdnn::wgtT_fc1, dnn::t5embdnn::bias_fc1);
379+
relu_activation<kHiddenFeatures>(h1);
380+
381+
linear_layer<kHiddenFeatures, kHiddenFeatures>(h1, h2, dnn::t5embdnn::wgtT_fc2, dnn::t5embdnn::bias_fc2);
382+
relu_activation<kHiddenFeatures>(h2);
383+
384+
linear_layer<kHiddenFeatures, Params_T5::kEmbed>(h2, embedding, dnn::t5embdnn::wgtT_fc3, dnn::t5embdnn::bias_fc3);
385+
}
386+
387+
} // namespace t5embdnn
388+
389+
namespace plsembdnn {
390+
template <typename TAcc>
391+
ALPAKA_FN_ACC ALPAKA_FN_INLINE void runEmbed(TAcc const& acc,
392+
const float eta,
393+
const float etaErr,
394+
const float phi,
395+
const float circleCenterX,
396+
const float circleCenterY,
397+
const float circleRadius,
398+
const float ptIn,
399+
const float ptErr,
400+
const bool isQuad,
401+
float (&embedding)[Params_pLS::kEmbed]) {
402+
constexpr unsigned int kInputFeatures = 10;
403+
constexpr unsigned int kHiddenFeatures = 32;
404+
405+
float x[kInputFeatures] = {eta / dnn::plsembdnn::kEta_norm,
406+
etaErr / dnn::plsembdnn::kEtaErr_norm,
407+
alpaka::math::cos(acc, phi),
408+
alpaka::math::sin(acc, phi),
409+
1.0f / ptIn,
410+
alpaka::math::log10(acc, ptErr),
411+
isQuad ? 1.0f : 0.0f,
412+
alpaka::math::log10(acc, alpaka::math::abs(acc, circleCenterX)),
413+
alpaka::math::log10(acc, alpaka::math::abs(acc, circleCenterY)),
414+
alpaka::math::log10(acc, circleRadius)};
415+
416+
float h1[kHiddenFeatures];
417+
float h2[kHiddenFeatures];
418+
419+
linear_layer<kInputFeatures, kHiddenFeatures>(x, h1, dnn::plsembdnn::wgtT_fc1, dnn::plsembdnn::bias_fc1);
420+
relu_activation<kHiddenFeatures>(h1);
421+
422+
linear_layer<kHiddenFeatures, kHiddenFeatures>(h1, h2, dnn::plsembdnn::wgtT_fc2, dnn::plsembdnn::bias_fc2);
423+
relu_activation<kHiddenFeatures>(h2);
424+
425+
linear_layer<kHiddenFeatures, Params_pLS::kEmbed>(
426+
h2, embedding, dnn::plsembdnn::wgtT_fc3, dnn::plsembdnn::bias_fc3);
427+
}
428+
429+
} // namespace plsembdnn
430+
291431
} // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst
292432

293433
#endif

0 commit comments

Comments
 (0)