Add settings for resource-efficient finetuning (#468)

Update resource settings
computational-cell-analytics · Mar 20, 2024 · 9f744e9 · 9f744e9
1 parent f746c5e
commit 9f744e9
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/micro_sam/training/training.py b/micro_sam/training/training.py
@@ -266,6 +266,10 @@ def default_sam_loader(**kwargs) -> DataLoader:
 # TODO extend this
 SETTINGS = {
     "Minimal": {"model_type": "vit_t", "n_objets_per_batch": 4, "n_sub_iterations": 4},
+    "CPU": {"model_type": "vit_b", "n_objects_per_batch": 10},
+    "gtx1080": {"model_type": "vit_t", "n_objects_per_batch": 5},
+    "rtx5000": {"model_type": "vit_b", "n_objects_per_batch": 10},
+    "V100": {"model_type": "vit_b"},
     "A100": {"model_type": "vit_h"},
 }