updates for v1.6 (#219)

* follow-up fixes to v1.5
jpata · Oct 2, 2023 · 6f0fbce · 6f0fbce
1 parent bab3205
commit 6f0fbce
Show file tree

Hide file tree

Showing 32 changed files with 317 additions and 189 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -31,7 +31,7 @@ jobs:
       - run: pip3 install torch==2.0.1 torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
       - run: pip install pyg-lib torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-2.0.1+cpu.html
 
-  tf-unittest:
+  tf-unittests:
     runs-on: ubuntu-22.04
     needs: [deps]
     steps:
@@ -41,12 +41,11 @@ jobs:
           python-version: '3.10'
           cache: 'pip'
       - run: pip install -r requirements.txt
-      - run: ./scripts/local_test_cms_pipeline.sh
       - run: PYTHONPATH=. python3 -m unittest tests/test_tf.py
 
   tf-clic-pipeline:
     runs-on: ubuntu-22.04
-    needs: [tf-unittest]
+    needs: [tf-unittests]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
@@ -58,7 +57,7 @@ jobs:
 
   tf-clic-hits-pipeline:
     runs-on: ubuntu-22.04
-    needs: [tf-unittest]
+    needs: [tf-unittests]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
@@ -70,7 +69,7 @@ jobs:
 
   tf-delphes-pipeline:
     runs-on: ubuntu-22.04
-    needs: [tf-unittest]
+    needs: [tf-unittests]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
@@ -82,7 +81,7 @@ jobs:
 
   tf-cms-pipeline:
     runs-on: ubuntu-22.04
-    needs: [tf-unittest]
+    needs: [tf-unittests]
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4

diff --git a/README_tf.md b/README_tf.md
@@ -2,17 +2,40 @@
 # Installation
 
 ```
-#get the code
 git clone https://github.com/jpata/particleflow.git
 cd particleflow
+git checkout v1.6
 ```
 
-## CLIC cluster-based training
+# Dependencies
 
+Install the required python packages, e.g.
 ```
-#Download the training datasets, about 5GB
+pip install -r requirements.txt
+```
+
+or use the singularity/apptainer environment for CUDA GPUs:
+```
+apptainer shell --nv https://hep.kbfi.ee/~joosep/tf-2.14.0.simg
+```
+
+# CLIC cluster-based training
+
+```
+#Download the training datasets, about 50GB
 rsync -r --progress lxplus.cern.ch:/eos/user/j/jpata/mlpf/tensorflow_datasets/clic/clusters/clic_* ~/tensorflow_datasets/
 
-#Run the training, multi-GPU support on the same machine is available, specify explicitly the GPUs you want to use
-CUDA_VISIBLE_DEVICES=0,1,2,3 python3 mlpf/pipeline.py train -c parameters/clic.yaml
+#Run the training from scratch, tuning the batch size to be suitable for an 8GB GPU
+python3 mlpf/pipeline.py train --config parameters/clic.yaml --batch-multiplier 0.5
+
+#Download and run the training from a checkpoint (optional)
+wget https://huggingface.co/jpata/particleflow/resolve/clic_clusters_v1.6/weights-96-5.346523.hdf5
+wget https://huggingface.co/jpata/particleflow/resolve/clic_clusters_v1.6/opt-96-5.346523.pkl
+python3 mlpf/pipeline.py train --config parameters/clic.yaml --weights weights-96-5.346523.hdf5 --batch-multiplier 0.5
+
+#Run the evaluation for a given training directory, loading the best weight file in the directory
+python3 mlpf/pipeline.py evaluate --train-dir experiments/clic-REPLACEME
+
+#Run the plots for a given training directory, using the evaluation outputs
+python3 mlpf/pipeline.py plots --train-dir experiments/clic-REPLACEME
 ```
diff --git a/mlpf/customizations.py b/mlpf/customizations.py
@@ -7,6 +7,9 @@ def customize_pipeline_test(config):
     # don't use dynamic batching, as that can result in weird stuff with very few events
     config["batching"]["bucket_by_sequence_length"] = False
 
+    if "delphes_pf_ttbar" in config["datasets"]:
+        config["train_test_datasets"]["physical"]["datasets"] = ["delphes_pf_ttbar"]
+
     # for cms, keep only ttbar
     if "cms_pf_ttbar" in config["datasets"]:
         config["train_test_datasets"]["physical"]["datasets"] = ["cms_pf_ttbar"]

diff --git a/mlpf/data_cms/README.md b/mlpf/data_cms/README.md
@@ -0,0 +1,10 @@
+```
+export IMG=/cvmfs/singularity.opensciencegrid.org/cmssw/cms:rhel7
+export SCRAM_ARCH=slc7_amd64_gcc10
+
+cmsrel CMSSW_12_3_0_pre6
+cd CMSSW_12_3_0_pre6
+git cms-init
+git-cms-merge-topic jpata:pfanalysis_caloparticle
+scram b -j4
+```
diff --git a/mlpf/data_cms/genjob.sh b/mlpf/data_cms/genjob.sh
@@ -1,8 +1,12 @@
 #!/bin/bash
+#SBATCH --partition main
+#SBATCH --cpus-per-task 1
+#SBATCH --mem-per-cpu 5G
+#SBATCH -o slurm-%x-%j-%N.out
 set -e
 set -x
 
-CMSSWDIR=/home/joosep/reco/mlpf/CMSSW_12_3_0_pre6
+CMSSWDIR=/home/joosep/CMSSW_12_3_0_pre6
 MLPF_PATH=/home/joosep/particleflow/
 
 #seed must be greater than 0

diff --git a/mlpf/data_cms/pu_files_local.txt b/mlpf/data_cms/pu_files_local.txt
@@ -0,0 +1,54 @@
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/00c14a87-07b5-4d7c-acb0-9c61d93677ea.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/0de932b0-f7f9-43d5-8a7f-68301921a476.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/12d51141-d984-4cc7-9ae0-425fc8c289bf.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/1d35d228-ac52-4c82-bd55-a5683673db94.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/22b726d9-b45f-4b5b-815f-bc095e0307c5.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/2f7e973a-c72d-417c-a2b0-c672281060b7.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/380b1531-aeaa-462c-a11a-0cc8e52a4d84.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/3b222f47-d811-43d3-9202-912a9c0230f7.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/3bbe5e05-77bb-4a4c-8e3e-fd742f36163b.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/3f755db8-f9f9-4978-8f42-a2da59a8f1a5.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/422c95e0-eb73-4da5-9069-b7cebc8c8cd5.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/42decb05-58f5-44cb-b081-4a996583d56d.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/474e9b8a-c5ed-469f-90ad-e424130cfa6b.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/4a0870ea-2b36-49e9-a20c-4833356b45ce.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/4b726c52-084b-47ea-9b73-893810c3ab7e.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/4d158869-649c-44ec-a214-be9cd54e3fcd.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/5ae5e6b9-5717-4d74-b29c-0494032f3884.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/5ca53d00-e129-4be8-a588-f3b80bc34e66.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/67e7ae25-4929-4d57-8cbe-427d36631015.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/6bac8d1f-7e0a-4eeb-8578-6fece23c5b8e.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/74132263-9d53-4c96-bbd0-bd14a19461a5.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/7f4440d2-1d7f-479c-be2a-5642ce7fea14.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/7f5a46c5-aa9c-420e-af61-13d124a8f703.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/8055fa51-0e9e-4dd0-ac1c-29b9afcd6b4a.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/807efaca-5d65-4589-9863-c52545360b86.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/81b2077f-3a06-416a-9d17-a090fffe2883.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/81ebcd02-6c69-47d2-a740-91ae93233924.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/85df0cb7-850a-44fd-8be5-bd4dfd482801.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/88006629-7690-45f0-876d-4017e8aee518.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/8915933b-9022-4b87-b6c4-c45b98c8cbb2.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/9693b840-70fb-43a8-9fd1-8a40c989dd47.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/9801beff-2c5a-449f-9458-153411a1619c.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/9f211a27-e26b-4db3-bee0-bc14310894aa.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/a28fd1e1-9799-46e2-b8de-5bdf5986f42b.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/a3bcd7b1-90fd-46eb-bc1e-042bac6ff938.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/a5456ebb-74e5-49df-9609-c04b2e4ef193.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/a7ca39f8-f64e-441c-9a12-6889c426b745.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/ad92595b-d421-4971-967c-a8124799ae73.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/b489ae40-5ece-479c-b3fa-5ee40ea1fa59.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/b6fdd1f7-55ea-46f2-afab-3af5636c1510.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/cb0c588c-982c-4d05-8f19-665705e17f06.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/d10a27be-7602-4d85-9356-fd1156823003.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/d960d8b1-8a20-4469-bd29-792f2a41e066.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/dcaaf8f9-bc0f-43ff-9eba-acf7f3dd75fd.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/e58a005f-389d-4f5f-afea-30eece093194.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/e5ec8cad-8eab-4b04-beba-ee1f92dc3896.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/ed9ac14a-168e-4e94-bf5c-801873dea749.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/f02197f1-b483-4806-a01f-62e8ef6a0009.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/f0b1a653-79f4-4653-8bae-0afe5d32ac68.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/f4e1f97b-f691-4ef8-a08e-624ae6d95062.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/f50c8bfe-41ca-406a-b2ab-21a9a09d0ac2.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/f5316554-e9bd-4380-a032-b556f756effa.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/fa1670e6-0248-488d-aa2a-1035ae71ba3c.root
+file:/store/relval/CMSSW_12_2_0_pre2/RelValMinBias_14TeV/GEN-SIM/122X_mcRun3_2021_realistic_v1_HighStat-v1/2580000/fd7520e5-eda6-47a0-8a3f-6b9c766b544b.root
diff --git a/mlpf/heptfds/delphes_pf/delphes_data_pf.py → mlpf/heptfds/delphes_pf/delphes_qcd_pf.py b/mlpf/heptfds/delphes_pf/delphes_data_pf.py → mlpf/heptfds/delphes_pf/delphes_qcd_pf.py
@@ -1,21 +1,24 @@
 from pathlib import Path
 
-from utils_delphes import prepare_data_delphes, X_FEATURES, Y_FEATURES
 import tensorflow_datasets as tfds
 import numpy as np
 
+from utils_delphes import X_FEATURES, Y_FEATURES
+from utils_delphes import split_sample, generate_examples
+
+
 _DESCRIPTION = """
 Dataset generated with Delphes.
 
-TTbar and QCD events with PU~200.
+QCD events with PU~200.
 """
 
 _CITATION = """
 https://zenodo.org/record/4559324#.YTs853tRVH4
 """
 
 
-class DelphesDataPf(tfds.core.GeneratorBasedBuilder):
+class DelphesQcdPf(tfds.core.GeneratorBasedBuilder):
     VERSION = tfds.core.Version("1.2.0")
     RELEASE_NOTES = {
         "1.0.0": "Initial release.",
@@ -28,7 +31,7 @@ class DelphesDataPf(tfds.core.GeneratorBasedBuilder):
 
     def __init__(self, *args, **kwargs):
         kwargs["file_format"] = tfds.core.FileFormat.ARRAY_RECORD
-        super(DelphesDataPf, self).__init__(*args, **kwargs)
+        super(DelphesQcdPf, self).__init__(*args, **kwargs)
 
     def _info(self) -> tfds.core.DatasetInfo:
         return tfds.core.DatasetInfo(
@@ -49,17 +52,7 @@ def _info(self) -> tfds.core.DatasetInfo:
 
     def _split_generators(self, dl_manager: tfds.download.DownloadManager):
         path = Path(dl_manager.manual_dir)
-        return {
-            "train": self._generate_examples(path / "pythia8_ttbar"),
-            "test": self._generate_examples(path / "pythia8_qcd"),
-        }
+        return split_sample(Path(path / "pythia8_qcd/raw"))
 
     def _generate_examples(self, path):
-        for fi in list(path.glob("*.pkl.bz2")):
-            Xs, ygens, ycands = prepare_data_delphes(str(fi))
-            for iev in range(len(Xs)):
-                yield str(fi) + "_" + str(iev), {
-                    "X": Xs[iev],
-                    "ygen": ygens[iev],
-                    "ycand": ycands[iev],
-                }
+        return generate_examples(path)
diff --git a/mlpf/heptfds/delphes_pf/delphes_ttbar_pf.py b/mlpf/heptfds/delphes_pf/delphes_ttbar_pf.py
@@ -0,0 +1,57 @@
+from pathlib import Path
+
+import tensorflow_datasets as tfds
+import numpy as np
+
+from utils_delphes import X_FEATURES, Y_FEATURES
+from utils_delphes import split_sample, generate_examples
+
+_DESCRIPTION = """
+Dataset generated with Delphes.
+
+TTbar events with PU~200.
+"""
+
+_CITATION = """
+https://zenodo.org/record/4559324#.YTs853tRVH4
+"""
+
+
+class DelphesTtbarPf(tfds.core.GeneratorBasedBuilder):
+    VERSION = tfds.core.Version("1.2.0")
+    RELEASE_NOTES = {
+        "1.0.0": "Initial release.",
+        "1.1.0": "Do not pad events to the same size",
+        "1.2.0": "Regenerate with ARRAY_RECORD",
+    }
+    MANUAL_DOWNLOAD_INSTRUCTIONS = """
+    Download from https://zenodo.org/record/4559324#.YTs853tRVH4
+    """
+
+    def __init__(self, *args, **kwargs):
+        kwargs["file_format"] = tfds.core.FileFormat.ARRAY_RECORD
+        super(DelphesTtbarPf, self).__init__(*args, **kwargs)
+
+    def _info(self) -> tfds.core.DatasetInfo:
+        return tfds.core.DatasetInfo(
+            builder=self,
+            description=_DESCRIPTION,
+            features=tfds.features.FeaturesDict(
+                {
+                    "X": tfds.features.Tensor(shape=(None, len(X_FEATURES)), dtype=np.float32),
+                    "ygen": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=np.float32),
+                    "ycand": tfds.features.Tensor(shape=(None, len(Y_FEATURES)), dtype=np.float32),
+                }
+            ),
+            supervised_keys=None,
+            homepage="https://zenodo.org/record/4559324#.YTs853tRVH4",
+            citation=_CITATION,
+            metadata=tfds.core.MetadataDict(x_features=X_FEATURES),
+        )
+
+    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+        path = Path(dl_manager.manual_dir)
+        return split_sample(Path(path / "pythia8_ttbar/raw"))
+
+    def _generate_examples(self, path):
+        return generate_examples(path)
diff --git a/mlpf/heptfds/delphes_pf/utils_delphes.py b/mlpf/heptfds/delphes_pf/utils_delphes.py
@@ -140,3 +140,30 @@ def prepare_data_delphes(fname, with_jet_idx=True):
         ycands.append(ycand)
 
     return Xs, ygens, ycands
+
+
+def split_sample(path, test_frac=0.8):
+    files = sorted(list(path.glob("*.pkl.bz2")))
+    print("Found {} files in {}".format(len(files), path))
+    assert len(files) > 0
+    idx_split = int(test_frac * len(files))
+    files_train = files[:idx_split]
+    files_test = files[idx_split:]
+    assert len(files_train) > 0
+    assert len(files_test) > 0
+    return {
+        "train": generate_examples(files_train),
+        "test": generate_examples(files_test),
+    }
+
+
+def generate_examples(files):
+    for fi in files:
+        Xs, ygens, ycands = prepare_data_delphes(str(fi))
+        assert len(Xs) > 0
+        for iev in range(len(Xs)):
+            yield str(fi) + "_" + str(iev), {
+                "X": Xs[iev],
+                "ygen": ygens[iev],
+                "ycand": ycands[iev],
+            }
diff --git a/mlpf/pipeline.py b/mlpf/pipeline.py
@@ -250,7 +250,11 @@ def train(
             pValue = ctypes.cast((ctypes.c_int * 1)(), ctypes.POINTER(ctypes.c_int))
             _libcudart.cudaDeviceSetLimit(ctypes.c_int(0x05), ctypes.c_int(128))
             _libcudart.cudaDeviceGetLimit(pValue, ctypes.c_int(0x05))
-            assert pValue.contents.value == 128
+            if pValue.contents.value != 128:
+                logging.warning(
+                    "Could not set L2 cache hint for GPU. Likely the GPU is too old for this optimization"
+                    + ", you may experience slower training performance."
+                )
 
     outdir = ""
     experiment = None

diff --git a/mlpf/plotting/plot_utils.py b/mlpf/plotting/plot_utils.py
@@ -104,7 +104,8 @@ def get_class_names(dataset_name):
 
 
 EVALUATION_DATASET_NAMES = {
-    "delphes_data_pf": r"Delphes-CMS $pp \rightarrow \mathrm{QCD}$",  # qcd is the validation set
+    "delphes_ttbar_pf": r"Delphes-CMS $pp \rightarrow \mathrm{t}\overline{\mathrm{t}}$",
+    "delphes_qcd_pf": r"Delphes-CMS $pp \rightarrow \mathrm{QCD}$",
     "cms_pf_qcd_high_pt": r"CMS high-$p_T$ QCD+PU events",
     "cms_pf_ttbar": r"CMS $\mathrm{t}\overline{\mathrm{t}}$+PU events",
     "cms_pf_single_neutron": r"CMS single neutron particle gun events",

diff --git a/mlpf/tfmodel/datasets/BaseDatasetFactory.py b/mlpf/tfmodel/datasets/BaseDatasetFactory.py
@@ -24,7 +24,10 @@ def unpack_target(y, num_output_classes, config):
     cos_phi = y[..., 5:6] * msk_pid
 
     type_as_int = tf.cast(y[..., 0], tf.int32)
-    charge_as_int = tf.cast(y[..., 1] + 1, tf.int32)  # -1, 0, 1 -> 0, 1, 2
+
+    # JP: in cms_pf_single_proton, charge is sometimes 2. so I added clip here currently
+    charge_as_int = tf.clip_by_value(tf.cast(y[..., 1] + 1, tf.int32), 0, 2)
+
     tf.debugging.assert_greater_equal(charge_as_int, 0, message="charge", summarize=100)
     tf.debugging.assert_less_equal(charge_as_int, 2, message="charge", summarize=100)