diff --git a/only_for_me/narval/train.py b/only_for_me/narval/train.py index 302a57a2..663b2880 100644 --- 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 index e7324790..791180c8 100644 --- 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