Skip to content

Commit

Permalink
update Sonic producer, customization functions
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpata
jpata committed Nov 6, 2020
1 parent 85e6622 commit 93abcd8
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 66 deletions.
130 changes: 69 additions & 61 deletions RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,21 @@
// git clone https://github.com/jpata/mlpf-models.git
// mv mlpf-models/triton/models ./
// singularity exec --nv docker://nvcr.io/nvidia/tritonserver:20.10-py3 tritonserver --model-repository models --log-verbose=1 --log-error=1 --log-info=1
// cmsRun src/RecoParticleFlow/PFProducer/test/mlpf_sonic.py
//
// process.mlpfproducergpu = cms.EDProducer("MLPFProducerSonicTriton",
// Client = cms.PSet(
// mode = cms.string("Sync"),
// batchSize = cms.untracked.uint32(1),
// address = cms.untracked.string("127.0.0.1"),
// port = cms.untracked.uint32(8001),
// timeout = cms.untracked.uint32(60),
// modelName = cms.string("mlpf"),
// modelVersion = cms.string("1"),
// verbose = cms.untracked.bool(False),
// allowedTries = cms.untracked.uint32(0),
// )
// )


class MLPFProducerSonicTriton : public SonicEDProducer<TritonClient> {
public:
Expand Down Expand Up @@ -104,76 +118,70 @@ class MLPFProducerSonicTriton : public SonicEDProducer<TritonClient> {
ielem += 1;
}
input1.toServer(data1);
edm::LogInfo("MLPFProducerSonic") << "acquire done";
}

void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override {
edm::LogInfo(client_.debugName()) << "produce";
using namespace reco::mlpf;

std::vector<reco::PFCandidate> pOutputCandidateCollection;
const auto& output1 = iOutput.at("Identity");
const auto& tmp = output1.template fromServer<float>();
for (const auto s : output1.shape()) {
edm::LogInfo(client_.debugName()) << "output shape=" << s;
}

//get the data of the first (and only) batch
const auto& out_data = output1.template fromServer<float>();

const auto num_elem = output1.shape()[1];

std::cout << "tmp.size" << tmp.size() << std::endl;
std::cout << "output.shape=" << output1.shape()[0] << "," << output1.shape()[1] << "," << output1.shape()[2]
<< std::endl;
for (long int ielem = 0; ielem < num_elem; ielem++) {
for (long int iprop = 0; iprop < 12; iprop++) {
std::cout << tmp[0][ielem * output1.shape()[2] + iprop] << " ";
unsigned int num_pred_particles = 0;
for (unsigned int ielem = 0; ielem < num_elem; ielem++) {
//get the coefficients in the output corresponding to the class probabilities (raw logits)
std::vector<float> pred_id_logits;
for (unsigned int idx_id = 0; idx_id <= NUM_CLASS; idx_id++) {
pred_id_logits.push_back(out_data[0][ielem*NUM_OUTPUTS + idx_id]);
}
std::cout << std::endl;
}

//edm::LogInfo(client_.debugName()) << "produced shape=" << tmp.size();
// edm::LogInfo(client_.debugName()) << "batchSize=" << client_.batchSize() << " noutput=" << client_.noutput();

// float output[client_.batchSize()][1000][NUM_OUTPUTS];

// for (unsigned int ibatch = 0; ibatch < client_.batchSize(); ibatch++) {
// for (unsigned int ielem = 0; ielem < 1000; ielem++) {
// for (unsigned int iprop = 0; iprop < NUM_OUTPUTS; iprop++) {
// output[ibatch][ielem][iprop] = iOutput.at(ibatch*1000 + ielem*NUM_OUTPUTS + iprop);
// }
// }
// }
// std::stringstream msg;
// for (unsigned int ibatch=0; ibatch<client_.batchSize(); ibatch++) {
// for (unsigned int ielem=0; ielem<1000; ielem++) {
// std::vector<float> pred_id_logits;
// for (unsigned int idx_id=0; idx_id <= NUM_CLASS; idx_id++) {
// pred_id_logits.push_back(output[ibatch][ielem][idx_id]);
// }
// //get the most probable class PDGID
// int pred_pid = pdgid_encoding.at(arg_max(pred_id_logits));
// if (pred_pid != 0) {
// float pred_eta = output[ibatch][ielem][IDX_ETA];
// float pred_phi = output[ibatch][ielem][IDX_PHI];
// pred_phi = atan2(sin(pred_phi), cos(pred_phi));
// float pred_e = output[ibatch][ielem][IDX_ENERGY];
// float pred_charge = output[ibatch][ielem][IDX_CHARGE];
// float pred_pt = pred_e / cosh(pred_eta);
// reco::PFCandidate::Charge charge = 0;
// if (pred_pid == 11 || pred_pid == 13 || pred_pid == 211) {
// charge = pred_charge > 0 ? +1 : -1;
// }
// TLorentzVector p4;
// p4.SetPtEtaPhiE(pred_pt, pred_eta, pred_phi, pred_e);

// reco::PFCandidate cand(0, math::XYZTLorentzVector(p4.X(), p4.Y(), p4.Z(), p4.E()), reco::PFCandidate::ParticleType(0));
// cand.setPdgId(pred_pid);
// cand.setCharge(charge);
// pOutputCandidateCollection.push_back(cand);
// }
// }
// }
// iEvent.emplace(pfCandidatesToken_, pOutputCandidateCollection);

//edm::LogInfo(client_.debugName()) << msg.str();
edm::LogInfo(client_.debugName()) << "produce done";
//get the most probable class PDGID
int pred_pid = pdgid_encoding[arg_max(pred_id_logits)];

//get the predicted momentum components
float pred_eta = out_data[0][ielem*NUM_OUTPUTS + IDX_ETA];
float pred_phi = out_data[0][ielem*NUM_OUTPUTS + IDX_PHI];

//put phi into a modular range
pred_phi = TMath::ATan2(TMath::Sin(pred_phi), TMath::Cos(pred_phi));
float pred_e = out_data[0][ielem*NUM_OUTPUTS + IDX_ENERGY];

//currently, set the pT from a massless approximation.
//later versions of the model may predict predict both the energy and pT of the particle
float pred_pt = pred_e / cosh(pred_eta);
float pred_charge = out_data[0][ielem*NUM_OUTPUTS + IDX_CHARGE];

//a particle was predicted for this PFElement
if (pred_pid != 0) {
//set the charge to +1 or -1 for PFCandidates that are charged, according to the sign of the predicted charge
reco::PFCandidate::Charge charge = 0;
if (pred_pid == 11 || pred_pid == 13 || pred_pid == 211) {
charge = pred_charge > 0 ? +1 : -1;
}

TLorentzVector p4;
p4.SetPtEtaPhiE(pred_pt, pred_eta, pred_phi, pred_e);

reco::PFCandidate cand(
0, math::XYZTLorentzVector(p4.X(), p4.Y(), p4.Z(), p4.E()), reco::PFCandidate::ParticleType(0));
cand.setPdgId(pred_pid);
cand.setCharge(charge);

pOutputCandidateCollection.push_back(cand);
num_pred_particles += 1;
} else {
//this element did not directly yield a particle, but may have been used indirectly for other ML-PFCandidates
//we can determine this from the sparse adjacency matrix that the model produces internally
}
} //loop over PFElements

iEvent.emplace(pfCandidatesToken_, pOutputCandidateCollection);
}

~MLPFProducerSonicTriton() override {}

//to ensure distinct cfi names - specialized below
Expand Down
32 changes: 27 additions & 5 deletions RecoParticleFlow/PFProducer/python/mlpfproducer_customise.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import FWCore.ParameterSet.Config as cms

def customise_step3(process, output="AODSIMoutput"):
def customise_step3(process, output="AODSIMoutput", usegpu=False):
if hasattr(process, "particleFlowBlock") and hasattr(process, output):
process.mlpfproducer = cms.EDProducer("MLPFProducer",
#we don't really need anything the block does, just to get all the PFElements in the event in a convenient list
Expand All @@ -16,11 +16,33 @@ def customise_step3(process, output="AODSIMoutput"):
out.outputCommands.append('keep *_ak4PFJets_*_*')
process.mlpf_path = cms.Path(process.mlpfproducer*process.ak4MLPFJets)

if usegpu:
process.mlpfproducergpu = cms.EDProducer("MLPFProducerSonicTriton",
Client = cms.PSet(
mode = cms.string("Sync"),
batchSize = cms.untracked.uint32(1),
address = cms.untracked.string("127.0.0.1"),
port = cms.untracked.uint32(8001),
timeout = cms.untracked.uint32(60),
modelName = cms.string("mlpf"),
modelVersion = cms.string("1"),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
)
)
process.mlpf_path.insert(-1, process.mlpfproducergpu)

process.schedule.insert(-1, process.mlpf_path)
return process

def customise_step3_reco(process, output="RECOSIMoutput"):
return customise_step3(process, output)
def customise_step3_reco(process):
return customise_step3(process, "RECOSIMoutput")

def customise_step3_aod(process):
return customise_step3(process, "AODSIMoutput")

def customise_step3_reco_gpu(process):
return customise_step3(process, "RECOSIMoutput", True)

def customise_step3_aod(process, output="AODSIMoutput"):
return customise_step3(process, output)
def customise_step3_aod_gpu(process):
return customise_step3(process, "AODSIMoutput", True)

0 comments on commit 93abcd8

Please sign in to comment.