1 gpu

only_for_me/narval/train.py

@@ -9,6 +9,8 @@
 from zoobot.pytorch.training import train_with_pytorch_lightning
 from zoobot.shared import benchmark_datasets, schemas
 
+import pytorch_lightning as pl
+
 
 if __name__ == '__main__':
 
@@ -47,6 +49,7 @@
     logging.basicConfig(level=logging.INFO)
 
     random_state = args.random_state
+    pl.seed_everything(random_state)
 
     # if args.nodes > 1:
     #     # at Manchester, our slurm cluster sets TASKS not NTASKS, which then confuses lightning
@@ -118,7 +121,7 @@
         resize_after_crop=args.resize_after_crop,
         # hardware parameters
         # gpus=args.gpus,
-        gpus=2,
+        gpus=1,
         nodes=args.nodes,
         mixed_precision=args.mixed_precision,
         wandb_logger=wandb_logger,

only_for_me/narval/train.sh

@@ -1,11 +1,11 @@
 #!/bin/bash
 #SBATCH --time=2:50:0  
 #SBATCH --nodes=1
-#SBATCH --ntasks=2
-#SBATCH --ntasks-per-node=2
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1
 #SBATCH --cpus-per-task=10
 #SBATCH --mem-per-cpu 4G
-#SBATCH --gres=gpu:v100:2
+#SBATCH --gres=gpu:v100:1
 
 nvidia-smi
 

zoobot/pytorch/training/debug_split.ipynb

@@ -126,6 +126,32 @@
     "wds.utils.pytorch_worker_info()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['5', '2']"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from itertools import islice\n",
+    "\n",
+    "def get_per_worker(urls={str(x) for x in range(10)}, worker_n=1, num_workers=5):\n",
+    "    for s in islice(urls, worker_n, None, num_workers):\n",
+    "        yield s\n",
+    "\n",
+    "list(get_per_worker(worker_n=0))"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,