From 101bf14cf277f72af15d1060ad37252da671d2d4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Th=C3=A9o?= Date: Mon, 8 Apr 2024 09:04:53 +0000 Subject: [PATCH] Have inferenceBatchSize as parameter --- .../TICL/plugins/TracksterLinkingbySuperClustering.cc | 9 +++++++-- .../TICL/plugins/TracksterLinkingbySuperClustering.h | 3 ++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.cc b/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.cc index 0e508781ec0d1..8c8698555168d 100644 --- a/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.cc +++ b/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.cc @@ -36,6 +36,7 @@ TracksterLinkingbySuperClustering::TracksterLinkingbySuperClustering(const edm:: cms::Ort::ONNXRuntime const* onnxRuntime) : TracksterLinkingAlgoBase(ps, iC, onnxRuntime), dnnVersion_(ps.getParameter("dnnVersion")), + inferenceBatchSize_(ps.getParameter("inferenceBatchSize")), nnWorkingPoint_(ps.getParameter("nnWorkingPoint")), deltaEtaWindow_(ps.getParameter("deltaEtaWindow")), deltaPhiWindow_(ps.getParameter("deltaPhiWindow")), @@ -101,9 +102,9 @@ void TracksterLinkingbySuperClustering::linkTracksters( /* Evaluate in minibatches since running with trackster count = 3000 leads to a short-lived ~15GB memory allocation Also we do not know in advance how many superclustering candidate pairs there are going to be - So we pick (arbitrarily) a batch size of 10^5. It needs to be rounded to featureCount + The batch size needs to be rounded to featureCount */ - const unsigned int miniBatchSize = static_cast(1e5) / nnInput->featureCount() * nnInput->featureCount(); + const unsigned int miniBatchSize = static_cast(inferenceBatchSize_) / nnInput->featureCount() * nnInput->featureCount(); std::vector> inputTensorBatches; // DNN input features tensors, in minibatches. Outer array : minibatches, inner array : 2D (flattened) array of features (indexed by batchIndex, featureId) @@ -318,6 +319,10 @@ void TracksterLinkingbySuperClustering::fillPSetDescription(edm::ParameterSetDes ->setComment("Path to DNN (as ONNX model)"); desc.add("dnnVersion", "alessandro-v2") ->setComment("DNN version tag. Can be alessandro-v1 or alessandro-v2"); + desc.add("inferenceBatchSize", 1e5) + ->setComment( + "Size of inference batches fed to DNN. Increasing it should produce faster inference but higher memory usage. " + "Has no physics impact."); desc.add("nnWorkingPoint", 0.51) ->setComment("Working point of DNN (in [0, 1]). DNN score above WP will attempt to supercluster."); desc.add("deltaEtaWindow", 0.1) diff --git a/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.h b/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.h index b0a35754f9996..d21737f664ab2 100644 --- a/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.h +++ b/RecoHGCal/TICL/plugins/TracksterLinkingbySuperClustering.h @@ -43,7 +43,8 @@ namespace ticl { private: bool checkExplainedVarianceRatioCut(ticl::Trackster const& ts) const; - const std::string dnnVersion_; // Version identifier of the DNN (to choose which inputs to use) + const std::string dnnVersion_; // Version identifier of the DNN (to choose which inputs to use) + unsigned int inferenceBatchSize_; // Size of inference batches fed to DNN double nnWorkingPoint_; // Working point for neural network (above this score, consider the trackster candidate for superclustering) float deltaEtaWindow_; // Delta eta window to consider trackster seed-candidate pairs for inference