Merge pull request #82 from ahoust17/master

Compatibility with Apple Silicon GPUs

atomai/__version__.py

@@ -1 +1 @@
-version = '0.7.8'
+version = '0.7.9'

atomai/trainers/trainer.py

@@ -67,7 +67,12 @@ class BaseTrainer:
  """
  def __init__(self):
  set_train_rng(1)
- self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ if torch.backends.mps.is_available():
+ self.device = torch.device("mps") # backend for Apple silicon GPUs
+ elif torch.cuda.is_available():
+ self.device = 'cuda'
+ else:
+ self.device = 'cpu'
  self.net = None
  self.criterion = None
  self.optimizer = None
@@ -362,6 +367,8 @@ def print_statistics(self, e: int, **kwargs) -> None:
  accuracy_metrics = "Accuracy"
  if torch.cuda.is_available():
  gpu_usage = gpu_usage_map(torch.cuda.current_device())
+ elif torch.backends.mps.is_available():
+ gpu_usage = gpu_usage_map(torch.mps.current_device())
  else:
  gpu_usage = ['N/A ', ' N/A']
  if self.compute_accuracy:

atomai/utils/nn.py

@@ -48,6 +48,8 @@ def load_weights(model: Type[torch.nn.Module],
  torch.manual_seed(0)
  if torch.cuda.device_count() > 0:
  checkpoint = torch.load(weights_path)
+ elif torch.backends.mps.is_available():
+ checkpoint = torch.load(weights_path, map_location='mps')
  else:
  checkpoint = torch.load(weights_path, map_location='cpu')
  model.load_state_dict(checkpoint)
@@ -199,6 +201,9 @@ def mock_forward(model: Type[torch.nn.Module],
  x = torch.randn(1, *dims)
  if next(model.parameters()).is_cuda:
  x = x.cuda()
+ elif next(model.parameters()).is_mps:
+ mps_device = torch.device("mps")
+ x = x.to(mps_device)
  out = model(x)
  return out