scaling of a100 2x repeat5x

mwalmsley · Nov 7, 2023 · 2d8c3bb · 2d8c3bb
1 parent 0f67ab6
commit 2d8c3bb
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/only_for_me/narval/train.py b/only_for_me/narval/train.py
@@ -121,7 +121,7 @@
         resize_after_crop=args.resize_after_crop,
         # hardware parameters
         # gpus=args.gpus,
-        gpus=1,
+        gpus=2,
         nodes=args.nodes,
         mixed_precision=args.mixed_precision,
         wandb_logger=wandb_logger,

diff --git a/only_for_me/narval/train.sh b/only_for_me/narval/train.sh
@@ -1,11 +1,11 @@
 #!/bin/bash
-#SBATCH --time=2:50:0  
+#SBATCH --time=1:00:0  
 #SBATCH --nodes=1
-#SBATCH --ntasks=1
-#SBATCH --ntasks-per-node=1
+#SBATCH --ntasks=2
+#SBATCH --ntasks-per-node=2
 #SBATCH --cpus-per-task=10
 #SBATCH --mem-per-cpu 4G
-#SBATCH --gres=gpu:v100:1
+#SBATCH --gres=gpu:a100:2
 
 nvidia-smi
 

diff --git a/zoobot/pytorch/training/webdatamodule.py b/zoobot/pytorch/training/webdatamodule.py
@@ -89,7 +89,7 @@ def make_loader(self, urls, mode="train"):
             # torch collate stacks dicts nicely while webdataset only lists them
             # so use the torch collate instead
             .batched(self.batch_size, torch.utils.data.default_collate, partial=False) 
-            # .repeat(5)
+            .repeat(5)
         )
 
         # from itertools import islice