Skip to content

Commit

Permalink
Timing the ONNX model, retrain CMS-GNNLSH-TF (jpata#229)
Browse files Browse the repository at this point in the history
  * timing code, retrain CMS-GNNLSH-TF, new multi-particle gun dataset, small hit-based datasets
  • Loading branch information
jpata authored and farakiko committed Jan 23, 2024
1 parent d456b00 commit 42929f6
Show file tree
Hide file tree
Showing 22 changed files with 723 additions and 176 deletions.
11 changes: 8 additions & 3 deletions mlpf/data_cms/genjob.sh
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,16 @@ MLPF_PATH=/home/joosep/particleflow/
SAMPLE=$1
SEED=$2

WORKDIR=`pwd`/$SAMPLE/$SEED
WORKDIR=/scratch/local/joosep/$SAMPLE/$SEED
#WORKDIR=`pwd`/$SAMPLE/$SEED
mkdir -p $WORKDIR

OUTDIR=/local/joosep/mlpf/cms/v2/$SAMPLE/raw
mkdir -p $OUTDIR

PILEUP=NoPileUp

N=1000
N=100

env
source /cvmfs/cms.cern.ch/cmsset_default.sh
Expand Down Expand Up @@ -73,4 +77,5 @@ cmsRun $CMSSWDIR/src/Validation/RecoParticleFlow/test/pfanalysis_ntuple.py
mv pfntuple.root pfntuple_${SEED}.root
python3 ${MLPF_PATH}/mlpf/data_cms/postprocessing2.py --input pfntuple_${SEED}.root --outpath ./ --save-normalized-table
bzip2 -z pfntuple_${SEED}.pkl
#rm step*.root
cp *.pkl.bz2 $OUTDIR/
rm -Rf $WORKDIR
46 changes: 22 additions & 24 deletions mlpf/data_cms/prepare_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,43 +3,41 @@

import os

outdir = "/local/joosep/mlpf/gen/v3/"
outdir = "/local/joosep/mlpf/cms/v2/"

samples = [
"SingleElectronFlatPt1To1000_pythia8_cfi",
"SingleGammaFlatPt1To1000_pythia8_cfi",
"SingleMuFlatPt1To1000_pythia8_cfi",
"SingleNeutronFlatPt0p7To1000_cfi",
"SinglePi0Pt1To1000_pythia8_cfi",
"SinglePiMinusFlatPt0p7To1000_cfi",
"SingleProtonMinusFlatPt0p7To1000_cfi",
"SingleTauFlatPt1To1000_cfi",
# "SingleElectronFlatPt1To1000_pythia8_cfi",
# "SingleGammaFlatPt1To1000_pythia8_cfi",
# "SingleMuFlatPt1To1000_pythia8_cfi",
# "SingleNeutronFlatPt0p7To1000_cfi",
# "SinglePi0Pt1To1000_pythia8_cfi",
# "SinglePiMinusFlatPt0p7To1000_cfi",
# "SingleProtonMinusFlatPt0p7To1000_cfi",
# "SingleTauFlatPt1To1000_cfi",
# "MultiParticlePFGun_cfi",
("MultiParticlePFGun50_cfi", 100000, 102050),
]

samples_pu = [
"TTbar_14TeV_TuneCUETP8M1_cfi",
"ZTT_All_hadronic_14TeV_TuneCUETP8M1_cfi",
"QCDForPF_14TeV_TuneCUETP8M1_cfi",
"QCD_Pt_3000_7000_14TeV_TuneCUETP8M1_cfi",
"SMS-T1tttt_mGl-1500_mLSP-100_TuneCP5_14TeV_pythia8_cfi",
"ZpTT_1500_14TeV_TuneCP5_cfi",
# "TTbar_14TeV_TuneCUETP8M1_cfi",
# "ZTT_All_hadronic_14TeV_TuneCUETP8M1_cfi",
# "QCDForPF_14TeV_TuneCUETP8M1_cfi",
# "QCD_Pt_3000_7000_14TeV_TuneCUETP8M1_cfi",
# "SMS-T1tttt_mGl-1500_mLSP-100_TuneCP5_14TeV_pythia8_cfi",
# "ZpTT_1500_14TeV_TuneCP5_cfi",
]

NUM_SAMPLES = 1050
SEED = 1

if __name__ == "__main__":

for s in samples_pu + samples:
for s, seed0, seed1 in samples_pu + samples:
is_pu = s in samples_pu

os.makedirs(outdir + "/" + s + "/raw", exist_ok=True)
os.makedirs(outdir + "/" + s + "/root", exist_ok=True)

for nsamp in range(NUM_SAMPLES):
if not os.path.isfile(outdir + "/" + s + "/raw/pfntuple_{}.pkl.bz2".format(SEED)):
for seed in range(seed0, seed1):
if not os.path.isfile(outdir + "/" + s + "/raw/pfntuple_{}.pkl.bz2".format(seed)):
if is_pu:
print("sbatch mlpf/tallinn/genjob_pu.sh {} {}".format(s, SEED))
print("sbatch genjob_pu.sh {} {}".format(s, seed))
else:
print("sbatch mlpf/tallinn/genjob.sh {} {}".format(s, SEED))
SEED += 1
print("sbatch genjob.sh {} {}".format(s, seed))
66 changes: 66 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/qq_10k.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample,
)

import tensorflow_datasets as tfds

from qq import _DESCRIPTION, _CITATION


class ClicEdmQqHitsPf10k(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.5.0")
RELEASE_NOTES = {
"1.5.0": "Regenerate with ARRAY_RECORD",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
For the raw input files in ROOT EDM4HEP format, please see the citation above.
The processed tensorflow_dataset can also be downloaded from:
FIXME
"""

def __init__(self, *args, **kwargs):
kwargs["file_format"] = tfds.core.FileFormat.ARRAY_RECORD
super(ClicEdmQqHitsPf10k, self).__init__(*args, **kwargs)

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample(Path(path / "p8_ee_qq_ecm380/"), max_files=100)

def _generate_examples(self, files):
return generate_examples(files)
66 changes: 66 additions & 0 deletions mlpf/heptfds/clic_pf_edm4hep_hits/ttbar_10k.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
from pathlib import Path

import tensorflow as tf
from utils_edm import (
X_FEATURES_CH,
X_FEATURES_TRK,
Y_FEATURES,
generate_examples,
split_sample,
)

import tensorflow_datasets as tfds

from ttbar import _DESCRIPTION, _CITATION


class ClicEdmTtbarHitsPf10k(tfds.core.GeneratorBasedBuilder):
VERSION = tfds.core.Version("1.5.0")
RELEASE_NOTES = {
"1.5.0": "Regenerate with ARRAY_RECORD",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
For the raw input files in ROOT EDM4HEP format, please see the citation above.
The processed tensorflow_dataset can also be downloaded from:
FIXME
"""

def __init__(self, *args, **kwargs):
kwargs["file_format"] = tfds.core.FileFormat.ARRAY_RECORD
super(ClicEdmTtbarHitsPf10k, self).__init__(*args, **kwargs)

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(
shape=(
None,
max(len(X_FEATURES_TRK), len(X_FEATURES_CH)),
),
dtype=tf.float32,
),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=None,
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(
x_features_track=X_FEATURES_TRK,
x_features_calohit=X_FEATURES_CH,
y_features=Y_FEATURES,
),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
path = dl_manager.manual_dir
return split_sample(Path(path / "p8_ee_tt_ecm380/"), max_files=100)

def _generate_examples(self, files):
return generate_examples(files)
4 changes: 3 additions & 1 deletion mlpf/heptfds/clic_pf_edm4hep_hits/utils_edm.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,8 +41,10 @@
labels = [0, 211, 130, 22, 11, 13]


def split_sample(path, test_frac=0.8):
def split_sample(path, test_frac=0.8, max_files=0):
files = sorted(list(path.glob("*.parquet")))
if max_files > 0:
files = files[:max_files]
print("Found {} files in {}".format(len(files), path))
assert len(files) > 0
idx_split = int(test_frac * len(files))
Expand Down
2 changes: 2 additions & 0 deletions mlpf/heptfds/cms_pf/cms_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,8 @@ def generate_examples(files):
x = Xs[ii]
yg = ygens[ii]
yc = ycands[ii]

uniqs, counts = np.unique(yg[:, 0], return_counts=True)
yield str(fi) + "_" + str(ii), {
"X": x,
"ygen": yg,
Expand Down
62 changes: 62 additions & 0 deletions mlpf/heptfds/cms_pf/multiparticlegun.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
"""CMS PF SinglePi dataset."""
import cms_utils
import tensorflow as tf

import tensorflow_datasets as tfds

X_FEATURES = cms_utils.X_FEATURES
Y_FEATURES = cms_utils.Y_FEATURES

_DESCRIPTION = """
Dataset generated with CMSSW and full detector sim.
Multi-particle gun events.
"""

# TODO(cms_pf): BibTeX citation
_CITATION = """
"""


class CmsPfMultiParticleGun(tfds.core.GeneratorBasedBuilder):
"""DatasetBuilder for cms_pf_multi_particle_gun dataset."""

VERSION = tfds.core.Version("1.6.0")
RELEASE_NOTES = {
"1.6.0": "Initial release",
}
MANUAL_DOWNLOAD_INSTRUCTIONS = """
rsync -r --progress lxplus.cern.ch:/eos/user/j/jpata/mlpf/cms/MultiParticlePFGun_cfi data/
"""

def __init__(self, *args, **kwargs):
kwargs["file_format"] = tfds.core.FileFormat.ARRAY_RECORD
super(CmsPfMultiParticleGun, self).__init__(*args, **kwargs)

def _info(self) -> tfds.core.DatasetInfo:
"""Returns the dataset metadata."""
# TODO(cms_pf): Specifies the tfds.core.DatasetInfo object
return tfds.core.DatasetInfo(
builder=self,
description=_DESCRIPTION,
features=tfds.features.FeaturesDict(
{
"X": tfds.features.Tensor(shape=(None, len(X_FEATURES)), dtype=tf.float32),
"ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
"ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=tf.float32),
}
),
supervised_keys=("X", "ycand"),
homepage="",
citation=_CITATION,
metadata=tfds.core.MetadataDict(x_features=X_FEATURES, y_features=Y_FEATURES),
)

def _split_generators(self, dl_manager: tfds.download.DownloadManager):
"""Returns SplitGenerators."""
path = dl_manager.manual_dir
sample_dir = "MultiParticlePFGun50_cfi"
return cms_utils.split_sample(path / sample_dir / "raw")

def _generate_examples(self, files):
return cms_utils.generate_examples(files)
56 changes: 0 additions & 56 deletions mlpf/tfmodel/timing.py

This file was deleted.

3 changes: 3 additions & 0 deletions mlpf/tfmodel/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -824,9 +824,12 @@ def model_scope(config, total_steps, weights=None, horovod_enabled=False):

def model_weight_setting():
grad_vars = model.trainable_weights
logging.info("grad_vars={}".format(len(grad_vars)))
zero_grads = [tf.zeros_like(w) for w in grad_vars]
logging.info("applying zero gradients to initialize optimizer")
opt.apply_gradients(zip(zero_grads, grad_vars))
if loaded_opt:
logging.info("setting optimizer state")
opt.set_weights(loaded_opt["weights"])

logging.info("distributing optimizer state")
Expand Down
Loading

0 comments on commit 42929f6

Please sign in to comment.