From beb801222f30f1c966e071bb43e24fdbf4add033 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Tue, 22 Mar 2022 13:05:24 +0100
Subject: [PATCH]  make samplers configurable

---
 template/ts-template.ipynb    |  9 ++++++---
 transformer/transformer.ipynb | 11 +++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/template/ts-template.ipynb b/template/ts-template.ipynb
index 5806c64..5f9b9d6 100644
--- a/template/ts-template.ipynb
+++ b/template/ts-template.ipynb
@@ -60,7 +60,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "8dc0844c",
    "metadata": {},
    "outputs": [],
@@ -99,6 +99,7 @@
     "    TestSplitSampler,\n",
     "    ExpectedNumInstanceSampler,\n",
     "    SelectFields,\n",
+    "    InstanceSampler,\n",
     ")\n",
     "from gluonts.torch.util import (\n",
     "    IterableDataset,\n",
@@ -345,6 +346,8 @@
     "        batch_size: int = 32,\n",
     "        num_batches_per_epoch: int = 50,\n",
     "        trainer_kwargs: Optional[Dict[str, Any]] = dict(),\n",
+    "        train_sampler: Optional[InstanceSampler] = None,\n",
+    "        validation_sampler: Optional[InstanceSampler] = None,\n",
     "    ) -> None:\n",
     "        trainer_kwargs = {\n",
     "            \"max_epochs\": 100,\n",
@@ -381,10 +384,10 @@
     "        self.batch_size = batch_size\n",
     "        self.num_batches_per_epoch = num_batches_per_epoch\n",
     "\n",
-    "        self.train_sampler = ExpectedNumInstanceSampler(\n",
+    "        self.train_sampler = train_sampler or ExpectedNumInstanceSampler(\n",
     "            num_instances=1.0, min_future=prediction_length\n",
     "        )\n",
-    "        self.validation_sampler = ValidationSplitSampler(\n",
+    "        self.validation_sampler = validation_sampler or ValidationSplitSampler(\n",
     "            min_future=prediction_length\n",
     "        )\n",
     "        \n",
diff --git a/transformer/transformer.ipynb b/transformer/transformer.ipynb
index 0af140d..1dec558 100644
--- a/transformer/transformer.ipynb
+++ b/transformer/transformer.ipynb
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "bc1a0f32",
    "metadata": {},
    "outputs": [],
@@ -56,6 +56,7 @@
     "    TestSplitSampler,\n",
     "    ExpectedNumInstanceSampler,\n",
     "    SelectFields,\n",
+    "    InstanceSampler,\n",
     ")\n",
     "from gluonts.torch.util import (\n",
     "    IterableDataset,\n",
@@ -570,6 +571,8 @@
     "        batch_size: int = 32,\n",
     "        num_batches_per_epoch: int = 50,\n",
     "        trainer_kwargs: Optional[Dict[str, Any]] = dict(),\n",
+    "        train_sampler: Optional[InstanceSampler] = None,\n",
+    "        validation_sampler: Optional[InstanceSampler] = None,\n",
     "    ) -> None:\n",
     "        trainer_kwargs = {\n",
     "            \"max_epochs\": 100,\n",
@@ -612,10 +615,10 @@
     "        self.batch_size = batch_size\n",
     "        self.num_batches_per_epoch = num_batches_per_epoch\n",
     "\n",
-    "        self.train_sampler = ExpectedNumInstanceSampler(\n",
+    "        self.train_sampler = train_sampler or ExpectedNumInstanceSampler(\n",
     "            num_instances=1.0, min_future=prediction_length\n",
     "        )\n",
-    "        self.validation_sampler = ValidationSplitSampler(\n",
+    "        self.validation_sampler = validation_sampler or ValidationSplitSampler(\n",
     "            min_future=prediction_length\n",
     "        )\n",
     "        \n",
@@ -1002,7 +1005,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.9.10"
   }
  },
  "nbformat": 4,