Timing the ONNX model, retrain CMS-GNNLSH-TF (jpata#229)

  * timing code, retrain CMS-GNNLSH-TF, new multi-particle gun dataset, small hit-based datasets

mlpf/data_cms/genjob.sh

@@ -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
@@ -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

mlpf/data_cms/prepare_args.py

@@ -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))

mlpf/heptfds/clic_pf_edm4hep_hits/qq_10k.py

@@ -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)

mlpf/heptfds/clic_pf_edm4hep_hits/ttbar_10k.py

@@ -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)

mlpf/heptfds/clic_pf_edm4hep_hits/utils_edm.py

@@ -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))

mlpf/heptfds/cms_pf/cms_utils.py

@@ -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,

mlpf/heptfds/cms_pf/multiparticlegun.py

@@ -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)

mlpf/tfmodel/utils.py

@@ -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")