From 693359c91042e208576f909a498ff138a85dc0d2 Mon Sep 17 00:00:00 2001 From: Genevieve Buckley <30920819+GenevieveBuckley@users.noreply.github.com> Date: Fri, 17 Jun 2022 11:23:19 -0400 Subject: [PATCH 1/2] Support pytorch acceleration on M1 mac hardware --- projects/mitonet/scripts/evaluate3d.py | 9 ++++++++- projects/mitonet/scripts/evaluate3d_bc.py | 9 ++++++++- .../mitonet/scripts/legacy_data/filter_nn.py | 20 ++++++++++++++++--- scripts/finetune.py | 17 +++++++++++++--- scripts/pdl_inference3d.py | 10 ++++++++-- 5 files changed, 55 insertions(+), 10 deletions(-) diff --git a/projects/mitonet/scripts/evaluate3d.py b/projects/mitonet/scripts/evaluate3d.py index 397e015..0260f85 100644 --- a/projects/mitonet/scripts/evaluate3d.py +++ b/projects/mitonet/scripts/evaluate3d.py @@ -71,7 +71,14 @@ def parse_args(): del state_dict[k] msg = model.load_state_dict(state['state_dict'], strict=True) - model.to('cuda' if torch.cuda.is_available() else 'cpu') # move model to GPU 0 + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + model.to(device) # set the evaluation transforms norms = state['norms'] diff --git a/projects/mitonet/scripts/evaluate3d_bc.py b/projects/mitonet/scripts/evaluate3d_bc.py index 3db0017..c0ef0dd 100644 --- a/projects/mitonet/scripts/evaluate3d_bc.py +++ b/projects/mitonet/scripts/evaluate3d_bc.py @@ -78,7 +78,14 @@ def parse_args(): del state_dict[k] msg = model.load_state_dict(state['state_dict'], strict=True) - model.to('cuda' if torch.cuda.device_count() > 0 else 'cpu') # move model to GPU 0 + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + model.to(device) # set the evaluation transforms norms = state['norms'] diff --git a/projects/mitonet/scripts/legacy_data/filter_nn.py b/projects/mitonet/scripts/legacy_data/filter_nn.py index 0323899..f99a887 100644 --- a/projects/mitonet/scripts/legacy_data/filter_nn.py +++ b/projects/mitonet/scripts/legacy_data/filter_nn.py @@ -58,7 +58,14 @@ # load the weights from online state_dict = torch.hub.load_state_dict_from_url(DEFAULT_WEIGHTS, map_location='cpu') msg = model.load_state_dict(state_dict) - model = model.to('cuda:0' if torch.cuda.is_available() else 'cpu') + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + model = model.to(device) model = model.eval() cudnn.benchmark = True @@ -97,8 +104,15 @@ def __getitem__(self, idx): tst_predictions = [] for data in tqdm(test, total=len(test)): with torch.no_grad(): - # load data onto gpu then forward pass - images = data['image'].to('cuda:0' if torch.cuda.is_available() else 'cpu', non_blocking=True) + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + # load data onto backend then do the forward pass + images = data['image'].to(device) output = model(images) predictions = nn.Sigmoid()(output) diff --git a/scripts/finetune.py b/scripts/finetune.py index f418df7..cf34d49 100644 --- a/scripts/finetune.py +++ b/scripts/finetune.py @@ -85,9 +85,20 @@ def main(): main_worker(config) def main_worker(config): - config['device'] = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - - if str(config['device']) == 'cpu': + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + config['device'] = device + + if str(config['device']) == 'cuda:0': + print("Using GPU for training.") + elif str(config['device']) == 'mps': + print("Using M1 Mac hardware for training.") + elif str(config['device']) == 'cpu': print(f"Using CPU for training.") # setup the model and pick dataset class diff --git a/scripts/pdl_inference3d.py b/scripts/pdl_inference3d.py index 8fbc475..03d5112 100644 --- a/scripts/pdl_inference3d.py +++ b/scripts/pdl_inference3d.py @@ -61,8 +61,14 @@ def parse_args(): # read the model config file config = load_config(args.config) - # set device and determine model to load - device = torch.device("cuda:0" if torch.cuda.is_available() and not args.use_cpu else "cpu") + # check whether GPU or M1 Mac hardware is available + if torch.cuda.is_available(): + device = torch.device('cuda:0') + elif torch.backends.mps.is_available(): + device = torch.device('mps') + else: + device = torch.device('cpu') + # determine model to load use_quantized = str(device) == 'cpu' and config.get('model_quantized') is not None model_key = 'model_quantized' if use_quantized else 'model' From 2cb93b049fe198eea9ae4bdc8a668ae28dda1b8e Mon Sep 17 00:00:00 2001 From: Genevieve Buckley <30920819+GenevieveBuckley@users.noreply.github.com> Date: Sun, 24 Jul 2022 15:24:33 +1000 Subject: [PATCH 2/2] Add comments explaining model types used --- empanada/models/panoptic_bifpn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/empanada/models/panoptic_bifpn.py b/empanada/models/panoptic_bifpn.py index b549b28..2b186e9 100644 --- a/empanada/models/panoptic_bifpn.py +++ b/empanada/models/panoptic_bifpn.py @@ -15,8 +15,8 @@ ) __all__ = [ - 'PanopticBiFPN', - 'PanopticBiFPNPR' + 'PanopticBiFPN', # simpler model that doesn't upsample. Uses PanopticDeepLabEngine (+ also won't upsample output automatically) + 'PanopticBiFPNPR' # what the empanada-napari plugin currently uses. Uses PanopticDeepLabRenderEngine ] class _BaseModel(nn.Module):