Merge pull request #48442 from SegmentLinking/reproducible_training

LST notebooks for neural network training: make training reproducible

Author: cmsbuild

RecoTracker/LSTCore/standalone/analysis/DNN/embed_train.ipynb

@@ -19,6 +19,13 @@
     "import awkward as ak # Using awkward array for easier handling of jagged data\n",
     "import time # For timing steps\n",
     "\n",
+    "# Set seeds for reproducibility\n",
+    "seed_value = 42\n",
+    "random.seed(seed_value)\n",
+    "np.random.seed(seed_value)\n",
+    "torch.manual_seed(seed_value)\n",
+    "os.environ['PYTHONHASHSEED'] = str(seed_value)\n",
+    "\n",
     "def load_root_file(file_path, branches=None, print_branches=False):\n",
     "    all_branches = {}\n",
     "    with uproot.open(file_path) as file:\n",
@@ -381,6 +388,7 @@
     "    dis_pairs = idxs_triu[dis_mask]\n",
     "\n",
     "    # down-sample\n",
+    "    random.seed(evt_idx)\n",
     "    if len(sim_pairs) > max_sim:\n",
     "        sim_pairs = sim_pairs[random.sample(range(len(sim_pairs)), max_sim)]\n",
     "    if len(dis_pairs) > max_dis:\n",
@@ -417,7 +425,7 @@
     "\n",
     "    with ProcessPoolExecutor(max_workers=n_workers) as pool:\n",
     "        futures = [pool.submit(_pairs_single_event, *args) for args in work_args]\n",
-    "        for fut in as_completed(futures):\n",
+    "        for fut in futures:\n",
     "            evt_idx, sim_pairs_evt, dis_pairs_evt = fut.result()\n",
     "            F = features_per_event[evt_idx]\n",
     "            D = displaced_per_event[evt_idx]\n",
@@ -546,6 +554,7 @@
     "    dis_pairs = np.column_stack((idx_p[dis_mask], idx_t[dis_mask]))\n",
     "\n",
     "    # down-sample\n",
+    "    random.seed(evt_idx)\n",
     "    if len(sim_pairs) > max_sim:\n",
     "        sim_pairs = sim_pairs[random.sample(range(len(sim_pairs)), max_sim)]\n",
     "    if len(dis_pairs) > max_dis:\n",
@@ -584,7 +593,7 @@
     "        )\n",
     "        for ev in range(len(features_per_event))\n",
     "    ]\n",
-    "    for fut in as_completed(futures):\n",
+    "    for fut in futures:\n",
     "        _, packed = fut.result()\n",
     "        # accumulate\n",
     "        sim_evt = sum(1 for _,_,lbl,_ in packed if lbl == 0)\n",

RecoTracker/LSTCore/standalone/analysis/DNN/train_T3_DNN.ipynb

@@ -1,5 +1,16 @@
 {
   "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# set seed for reproducibility\n",
+        "import torch\n",
+        "torch.manual_seed(42)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 1,

RecoTracker/LSTCore/standalone/analysis/DNN/train_T5_DNN.ipynb

@@ -1,5 +1,16 @@
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# set seed for reproducibility\n",
+    "import torch\n",
+    "torch.manual_seed(42)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 1,

RecoTracker/LSTCore/standalone/analysis/DNN/train_pT3_DNN.ipynb

@@ -1,5 +1,16 @@
 {
   "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# set seed for reproducibility\n",
+        "import torch\n",
+        "torch.manual_seed(42)"
+      ]
+    },
     {
       "cell_type": "code",
       "execution_count": 1,