feat: ✨ Add support for training on Apple M1/M2/M3 (mps) devices.

Also add log printing of device used for training.
janelia-cellmap · Oct 24, 2024 · bb9268b · bb9268b
1 parent 2d11926
commit bb9268b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/dacapo/compute_context/local_torch.py b/dacapo/compute_context/local_torch.py
@@ -60,6 +60,8 @@ def device(self):
                 if free < self.oom_limit:  # less than 1 GB free, decrease chance of OOM
                     return torch.device("cpu")
                 return torch.device("cuda")
+            elif torch.backends.mps.is_available():
+                return torch.device("mps")
             else:
                 return torch.device("cpu")
         return torch.device(self._device)
diff --git a/dacapo/train.py b/dacapo/train.py
@@ -135,6 +135,7 @@ def train_run(run: Run, do_validate=True):
     compute_context = create_compute_context()
     run.model = run.model.to(compute_context.device)
     run.move_optimizer(compute_context.device)
+    logger.info(f"Training on {compute_context.device}")
 
     array_store = create_array_store()
     run.trainer.iteration = trained_until