End_to_End_Diarization_Training.ipynb (#14680)

tango4j · web-flow · commit 57ef50d74f9f · 2025-09-08T19:03:49.000-04:00
Signed-off-by: taejinp &lt;tango4j@gmail.com&gt;
diff --git a/tutorials/speaker_tasks/End_to_End_Diarization_Training.ipynb b/tutorials/speaker_tasks/End_to_End_Diarization_Training.ipynb
@@ -778,7 +778,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Train a Sortformer Diarizer Model"
+    "## Train a Sortformer Diarizer Model\n",
+    "\n",
+    "### Training an offline Sortformer model"
    ]
   },
   {
@@ -823,7 +825,7 @@
    "source": [
     "curr_dir = os.getcwd() + \"/\"\n",
     "config.model.train_ds.manifest_filepath = f'{curr_dir}simulated_train/sortformer_train.json'\n",
-    "config.model.test_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
+    "# config.model.test_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
     "config.model.validation_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
     "config.trainer.strategy = \"ddp_notebook\"\n",
     "config.batch_size = 3\n",
@@ -858,11 +860,68 @@
     "sortformer_model.maybe_init_from_pretrained_checkpoint(config)\n",
     "trainer.fit(sortformer_model)"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Training a streaming Sortformer model\n",
+    "\n",
+    "If you want to train a streaming version of Sortformer, you can download the following YAML file."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "!wget -P conf https://raw.githubusercontent.com/NVIDIA/NeMo/{BRANCH}/examples/speaker_tasks/diarization/conf/neural_diarizer/streaming_sortformer_diarizer_4spk-v2.yaml\n",
+    "MODEL_CONFIG = os.path.join(NEMO_ROOT,'conf/streaming_sortformer_diarizer_4spk-v2.yaml')\n",
+    "config = OmegaConf.load(MODEL_CONFIG)\n",
+    "\n",
+    "curr_dir = os.getcwd() + \"/\"\n",
+    "config.model.train_ds.manifest_filepath = f'{curr_dir}simulated_train/sortformer_train.json'\n",
+    "config.model.test_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
+    "config.model.validation_ds.manifest_filepath = f'{curr_dir}simulated_valid/sortformer_valid.json'\n",
+    "config.trainer.strategy = \"ddp_notebook\"\n",
+    "config.batch_size = 3\n",
+    "\n",
+    "config.trainer.devices=1\n",
+    "config.accelerator=\"gpu\"\n",
+    "print(os.getcwd())\n",
+    "\n",
+    "print(\"config.model.train_ds.manifest_filepath \", config.model.train_ds.manifest_filepath )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Initiate a streaming Sortformer diarization training session using the given configurations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(devices=1, accelerator='gpu', max_epochs=50,\n",
+    "                  enable_checkpointing=False, logger=False,\n",
+    "                  log_every_n_steps=5, check_val_every_n_epoch=10)\n",
+    "\n",
+    "exp_manager(trainer, config.get(\"exp_manager\", None))\n",
+    "streaming_sortformer_model = SortformerEncLabelModel(cfg=config.model, trainer=trainer)\n",
+    "streaming_sortformer_model.maybe_init_from_pretrained_checkpoint(config)\n",
+    "trainer.fit(streaming_sortformer_model)"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "nv082124",
    "language": "python",
    "name": "python3"
   },
