Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Introduce Unified Particle Transformer v2 #47173

Merged
merged 8 commits into from
Feb 4, 2025
Merged

Introduce Unified Particle Transformer v2 #47173

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
2333867
UParT v2 update
AlexDeMoor Jan 23, 2025
e684e8c
UParT taginfo update for v2
AlexDeMoor Jan 23, 2025
ce85d68
code format
AlexDeMoor Jan 23, 2025
0e8a812
Update code (dir naming + clamp the const num)
AlexDeMoor Jan 29, 2025
4c3ec0d
Add UParT_v1 regression
Feb 1, 2025
570033a
Remove useless onnx model from CMSSW
Feb 1, 2025
45da89d
Using the same addUnifiedParTAK4 in NanoAOD for both UParT v1 and v2
Feb 3, 2025
eff5b69
Update UParT v1 jet neut regression description in Nano
AlexDeMoor Feb 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion PhysicsTools/NanoAOD/python/custom_jme_cff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,10 +27,11 @@
from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import _pfParticleTransformerAK4JetTagsAll
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll
bTagDiscriminatorsForAK4 = cms.PSet(foo = cms.vstring(
bTagDeepJet+
_pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll+_pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll+
_pfParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4JetTagsAll
_pfParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4JetTagsAll + _pfUnifiedParticleTransformerAK4V1JetTagsAll
))
run2_nanoAOD_ANY.toModify(
bTagDiscriminatorsForAK4,
Expand Down
7 changes: 7 additions & 0 deletions PhysicsTools/NanoAOD/python/jetsAK4_Puppi_cff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,9 @@
UParTAK4RegPtRawCorr = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptcorr'):-1",float,precision=10,doc="UnifiedParT universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT"),
UParTAK4RegPtRawCorrNeutrino = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu')>0?bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptnu'):-1",float,precision=10,doc="UnifiedParT universal flavor-aware pT regression neutrino correction, relative to visible. Correction relative to raw jet pT"),
UParTAK4RegPtRawRes = Var("?(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')+bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow'))>0?0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4JetTags:ptreslow')):-1",float,precision=10,doc="UnifiedParT universal flavor-aware jet pT resolution estimator, (q84 - q16)/2"),
UParTAK4V1RegPtRawCorr = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptcorr')>0?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptcorr'):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware visible pT regression (no neutrinos), correction relative to raw jet pT"),
UParTAK4V1RegPtRawCorrNeutrino = Var("?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptnu')>0?bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptnu'):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware pT regression neutrino correction, relative to visible. Correction relative to raw jet pT"),
UParTAK4V1RegPtRawRes = Var("?(bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreshigh')+bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreslow'))>0?0.5*(bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreshigh')-bDiscriminator('pfUnifiedParticleTransformerAK4V1JetTags:ptreslow')):-1",float,precision=10,doc="UnifiedParT V1 universal flavor-aware jet pT resolution estimator, (q84 - q16)/2"),
puIdDisc = Var("userFloat('pileupJetIdPuppi:fullDiscriminant')", float,doc="Pileup ID BDT discriminant with 133X Winter24 PuppiV18 training",precision=10),
hfsigmaEtaEta = Var("userFloat('hfsigmaEtaEta')",float,doc="sigmaEtaEta for HF jets (noise discriminating variable)",precision=10),
hfsigmaPhiPhi = Var("userFloat('hfsigmaPhiPhi')",float,doc="sigmaPhiPhi for HF jets (noise discriminating variable)",precision=10),
Expand Down Expand Up @@ -151,6 +154,9 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
print("Updating process to run UnifiedParTAK4")
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll as pfUnifiedParticleTransformerAK4JetTagsAll
_btagDiscriminators += pfUnifiedParticleTransformerAK4JetTagsAll
print("Updating process to run UnifiedParTAK4V1")
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll as pfUnifiedParticleTransformerAK4V1JetTagsAll
_btagDiscriminators += pfUnifiedParticleTransformerAK4V1JetTagsAll
if len(_btagDiscriminators)==0: return process
print("Will recalculate the following discriminators: "+", ".join(_btagDiscriminators))
updateJetCollection(
Expand All @@ -166,6 +172,7 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
return process

nanoAOD_addDeepInfoAK4_switch = cms.PSet(

nanoAOD_addParticleNet_switch = cms.untracked.bool(False),
nanoAOD_addRobustParTAK4Tag_switch = cms.untracked.bool(False),
nanoAOD_addUnifiedParTAK4Tag_switch = cms.untracked.bool(False)
Expand Down
13 changes: 13 additions & 0 deletions PhysicsTools/PatAlgos/python/recoLayer0/bTagging_cff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@
, 'pfParticleTransformerAK4TagInfos'
# UnifiedParticleTransformerAK4 tag infos
, 'pfUnifiedParticleTransformerAK4TagInfos'
# UnifiedParticleTransformerAK4V1 tag infos
, 'pfUnifiedParticleTransformerAK4V1TagInfos'
# GlobalParticleTransformerAK8 tag infos
, 'pfGlobalParticleTransformerAK8TagInfos'
# DeepDoubleB/C tag infos
Expand Down Expand Up @@ -367,6 +369,17 @@
supportedMetaDiscr[disc] = _pfUnifiedParticleTransformerAK4JetTagsProbs
# -----------------------------------

# -----------------------------------
# setup UnifiedParticleTransformer AK4 V1
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsProbs, _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs
# update supportedBtagDiscr
for disc in _pfUnifiedParticleTransformerAK4V1JetTagsProbs + _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs:
supportedBtagDiscr[disc] = [["pfUnifiedParticleTransformerAK4V1TagInfos"]]
# update supportedMetaDiscr
for disc in _pfUnifiedParticleTransformerAK4V1JetTagsMetaDiscrs:
supportedMetaDiscr[disc] = _pfUnifiedParticleTransformerAK4V1JetTagsProbs
# -----------------------------------

# -----------------------------------
# setup Negative ParticleTransformer AK4
from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import _pfNegativeParticleTransformerAK4JetTagsProbs
Expand Down
4 changes: 3 additions & 1 deletion PhysicsTools/PatAlgos/python/slimming/applyDeepBtagging_cff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@ def applyDeepBtagging(process, postfix=""):
from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll as pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4_cff import _pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll as pfParticleNetFromMiniAODAK4CHSCentralJetTagsAll
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import _pfUnifiedParticleTransformerAK4JetTagsAll as pfUnifiedParticleTransformerAK4JetTagsAll

from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import _pfUnifiedParticleTransformerAK4V1JetTagsAll as pfUnifiedParticleTransformerAK4V1JetTagsAll

# update slimmed jets to include DeepFlavour (keep same name)
# make clone for DeepFlavour-less slimmed jets, so output name is preserved
addToProcessAndTask('slimmedJetsNoDeepFlavour', slimmedJets.clone(), process, task)
Expand Down Expand Up @@ -56,6 +57,7 @@ def applyDeepBtagging(process, postfix=""):
+ pfParticleNetFromMiniAODAK4PuppiCentralJetTagsAll
+ pfParticleNetFromMiniAODAK4PuppiForwardJetTagsAll
+ pfUnifiedParticleTransformerAK4JetTagsAll
+ pfUnifiedParticleTransformerAK4V1JetTagsAll
)

updateJetCollection(
Expand Down
31 changes: 31 additions & 0 deletions PhysicsTools/PatAlgos/python/tools/jetTools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,6 +730,37 @@ def setupBTagging(process, jetSource, pfCandidates, explicitJTA, pvSource, svSou
flip = flip),
process, task)

if 'UnifiedParticleTransformerAK4V1TagInfos' in btagInfo:
svUsed = svSource
if btagInfo == 'pfNegativeUnifiedParticleTransformerAK4V1TagInfos':
svUsed = cms.InputTag(btagPrefix+'inclusiveCandidateNegativeSecondaryVertices'+labelName+postfix)
flip = True
else:
flip = False
# use right input tags when running with RECO PF candidates, which actually
# depends of whether jets use "particleFlow"
if pfCandidates.value() == 'packedPFCandidates':
puppi_value_map = setupPuppiForPackedPF(process)[0]
vertex_associator = cms.InputTag("")
else:
puppi_value_map = cms.InputTag("puppi")
vertex_associator = cms.InputTag("primaryVertexAssociation","original")

# If this jet is a puppi jet, then set is_weighted_jet to true.
is_weighted_jet = False
if ('puppi' in jetSource.value().lower()):
is_weighted_jet = True
addToProcessAndTask(btagPrefix+btagInfo+labelName+postfix,
btag.pfUnifiedParticleTransformerAK4V1TagInfos.clone(
jets = jetSource,
vertices=pvSource,
secondary_vertices=svUsed,
puppi_value_map = puppi_value_map,
vertex_associator = vertex_associator,
is_weighted_jet = is_weighted_jet,
flip = flip),
process, task)

if btagInfo == 'pfDeepDoubleXTagInfos':
# can only run on PAT jets, so the updater needs to be used
if 'updated' not in jetSource.value().lower():
Expand Down
1 change: 1 addition & 0 deletions RecoBTag/Configuration/python/RecoBTag_cff.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
from RecoBTag.ONNXRuntime.pfParticleNetAK4_cff import *
from RecoBTag.ONNXRuntime.pfParticleTransformerAK4_cff import *
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4_cff import *
from RecoBTag.ONNXRuntime.pfUnifiedParticleTransformerAK4V1_cff import *
from RecoBTag.ONNXRuntime.pfGlobalParticleTransformerAK8_cff import *
from RecoVertex.AdaptiveVertexFinder.inclusiveVertexing_cff import *
from RecoBTag.PixelCluster.pixelClusterTagInfos_cfi import *
Expand Down
2 changes: 1 addition & 1 deletion RecoBTag/FeatureTools/plugins/UnifiedParticleTransformerAK4TagInfoProducer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,7 +159,7 @@ void UnifiedParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::Configu
desc.add<double>("min_candidate_pt", 0.10);
desc.add<bool>("flip", false);
desc.add<bool>("sort_cand_by_pt", false);
desc.add<bool>("fix_lt_sorting", false);
desc.add<bool>("fix_lt_sorting", true);
desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
desc.add<edm::InputTag>("losttracks", edm::InputTag("lostTracks"));
desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
Expand Down
25 changes: 19 additions & 6 deletions RecoBTag/ONNXRuntime/plugins/UnifiedParticleTransformerAK4ONNXJetTagsProducer.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ class UnifiedParticleTransformerAK4ONNXJetTagsProducer : public edm::stream::EDP
const edm::EDGetTokenT<TagInfoCollection> src_;
std::vector<std::string> flav_names_;
std::vector<std::string> input_names_;
bool use_dynamic_axes_ = false;
std::vector<std::string> output_names_;

enum InputIndexes {
Expand Down Expand Up @@ -76,6 +77,8 @@ UnifiedParticleTransformerAK4ONNXJetTagsProducer::UnifiedParticleTransformerAK4O
: src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
use_dynamic_axes_(iConfig.getParameter<edm::FileInPath>("model_path").fullPath().find("v2.onnx") !=
std::string::npos),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would suggest adding a configurable parameter for use_dynamic_axes (with a default value to true), otherwise we are stuck with the naming of v2.onnx.

output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
// get output names from flav_names
for (const auto& flav_name : flav_names_) {
Expand All @@ -89,7 +92,8 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::fillDescriptions(edm::Con
desc.add<edm::InputTag>("src", edm::InputTag("pfUnifiedParticleTransformerAK4TagInfos"));
desc.add<std::vector<std::string>>(
"input_names", {"input_1", "input_2", "input_3", "input_4", "input_5", "input_6", "input_7", "input_8"});
desc.add<edm::FileInPath>("model_path", edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V00/UParTAK4.onnx"));
desc.add<edm::FileInPath>("model_path",
edm::FileInPath("RecoBTag/Combined/data/UParTAK4/PUPPI/V01/UParTAK4_v2.onnx"));
desc.add<std::vector<std::string>>("output_names", {"softmax"});
desc.add<std::vector<std::string>>(
"flav_names",
Expand Down Expand Up @@ -164,11 +168,20 @@ void UnifiedParticleTransformerAK4ONNXJetTagsProducer::get_input_sizes(
const reco::FeaturesTagInfo<btagbtvdeep::UnifiedParticleTransformerAK4Features> taginfo) {
const auto& features = taginfo.features();

/// We require a fixed size due to an ONNX conversion issue (to be improved in the future ?) ///
n_cpf_ = (unsigned int)29;
n_lt_ = (unsigned int)5;
n_npf_ = (unsigned int)25;
n_sv_ = (unsigned int)5;
if (use_dynamic_axes_) {
// Use actual sizes for dynamic axes version
n_cpf_ = std::clamp((unsigned int)features.c_pf_features.size(), (unsigned int)1, (unsigned int)29);
n_lt_ = std::clamp((unsigned int)features.lt_features.size(), (unsigned int)1, (unsigned int)5);
n_npf_ = std::clamp((unsigned int)features.n_pf_features.size(), (unsigned int)1, (unsigned int)25);
n_sv_ = std::clamp((unsigned int)features.sv_features.size(), (unsigned int)1, (unsigned int)5);

} else {
// Use fixed sizes for original version
n_cpf_ = (unsigned int)29;
n_lt_ = (unsigned int)5;
n_npf_ = (unsigned int)25;
n_sv_ = (unsigned int)5;
}

input_sizes_ = {
n_cpf_ * n_features_cpf_,
Expand Down
Loading