prepare for terrestrial

mwalmsley · Nov 30, 2023 · 0c1d6e2 · 0c1d6e2
1 parent fda028a
commit 0c1d6e2
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 10 deletions.
diff --git a/only_for_me/narval/train.py b/only_for_me/narval/train.py
@@ -25,6 +25,8 @@
     # parser.add_argument('--data-dir', dest='data_dir', type=str)
     # parser.add_argument('--dataset', dest='dataset', type=str, help='dataset to use, either "gz_decals_dr5" or "gz_evo"')
     parser.add_argument('--architecture', dest='architecture_name', default='efficientnet_b0', type=str)
+    parser.add_argument('--accumulate-gradients', dest='acculumate_gradients', default=1, type=int)
+    parser.add_argument('--terrestrial-init', dest='terrestrial', default=False, action='store_true')
     parser.add_argument('--resize-after-crop', dest='resize_after_crop',
                         type=int, default=224)
     parser.add_argument('--color', default=False, action='store_true')
@@ -120,6 +122,9 @@
     if args.num_features != 1280:
         timm_kwargs.update({'num_features': args.num_features})
 
+    if args.terrestrial:
+        timm_kwargs.update({'pretrained': True})
+
     train_with_pytorch_lightning.train_default_zoobot_from_scratch(
         save_dir=args.save_dir,
         schema=schema,

diff --git a/only_for_me/narval/train.sh b/only_for_me/narval/train.sh
@@ -23,7 +23,16 @@ REPO_DIR=/project/def-bovy/walml/zoobot
 #     --num-features 128 \
 #     --gpus 1 \
 #     --num-workers 10 \
-#     --color --wandb --mixed-precision --compile-encoder
+#     --color --wandb --mixed-precision 
+
+srun $PYTHON $REPO_DIR/only_for_me/narval/train.py \
+    --save-dir $REPO_DIR/only_for_me/narval/desi_300px_maxvittiny_1gpu \
+    --batch-size 64 \
+    --gpus 1 \
+    --num-workers 10 \
+    --color --wandb --mixed-precision 
+
+    # \ --compile-encoder
 
 
 # batch sizes
@@ -47,15 +56,15 @@ REPO_DIR=/project/def-bovy/walml/zoobot
 # efficientnet_b5 - 64. remember it expects bigger images tho, may not work great
 # maxvit_rmlp_base_rw_224 - 32 (95%). Now scaling at 16 gpus
 
-srun $PYTHON $REPO_DIR/only_for_me/narval/train.py \
-    --save-dir $REPO_DIR/only_for_me/narval/desi_300px_maxvit_rmlp_base_rw_224_4gpu_w005 \
-    --batch-size 32 \
-    --gpus 4 \
-    --nodes 1 \
-    --num-workers 5 \
-    --weight-decay 0.05 \
-    --architecture maxvit_rmlp_base_rw_224 \
-    --color --wandb --mixed-precision
+# srun $PYTHON $REPO_DIR/only_for_me/narval/train.py \
+#     --save-dir $REPO_DIR/only_for_me/narval/desi_300px_maxvit_rmlp_base_rw_224_4gpu_w005 \
+#     --batch-size 32 \
+#     --gpus 4 \
+#     --nodes 1 \
+#     --num-workers 5 \
+#     --weight-decay 0.05 \
+#     --architecture maxvit_rmlp_base_rw_224 \
+#     --color --wandb --mixed-precision
 
     #  --compile-encoder