Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for M1/M2 mac in PytorchModelHandlerTensor #29999

Closed
wants to merge 3 commits into from
Closed

Add support for M1/M2 mac in PytorchModelHandlerTensor #29999

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
255656a
Update pytorch_inference.py
bfcrampton Jan 12, 2024
416c3ad
Support both Pytorch classes
bfcrampton Jan 12, 2024
55db8bb
Address comments
bfcrampton Jan 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 20 additions & 6 deletions sdks/python/apache_beam/ml/inference/pytorch_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,12 @@ def _load_model(
"Model handler specified a 'GPU' device, but GPUs are not available. "
"Switching to CPU.")
device = torch.device('cpu')
if device == torch.device('mps') and not (torch.backend.mps.is_available() and
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, it also looks like torch.device('mps') causes older versions of pytorch to throw. To maintain compat with older versions, could we wrap this in a try except? It should be fine to silently swallow the exception since we'll throw below if someone tries to set the device to mps and the pytorch version doesn't support it.

This is why the test checks are failing

torch.backend.mps.is_built()):
logging.warning(
"Model handler specified a 'MPS' device, but it is not available. "
"Switching to CPU.")
device = torch.device('cpu')

try:
logging.info(
Expand Down Expand Up @@ -216,8 +222,9 @@ def __init__(
model_params: A dictionary of arguments required to instantiate the model
class.
device: the device on which you wish to run the model. If
``device = GPU`` then a GPU device will be used if it is available.
Otherwise, it will be CPU.
``device = GPU`` then a cuda device will be used if it is available.
If ``device = MPS`` then a mps device will be used if it is available.
Otherwise, it will be cpu.
inference_fn: the inference function to use during RunInference.
default=_default_tensor_inference_fn
torch_script_model_path: Path to the torch script model.
Expand All @@ -243,9 +250,12 @@ def __init__(
with PyTorch 1.9 and 1.10.
"""
self._state_dict_path = state_dict_path
if device == 'GPU':
if device in ['gpu', 'GPU', 'cuda', 'CUDA']:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey, sorry for the slow review - this looks mostly good. The only thing I'd ask is that if we go case-insensitive we do it all the way (so device.lower() == 'gpu' here, and the same type of change applied elsewhere in the PR). Otherwise, this looks good to me, thanks!

logging.info("Device is set to CUDA")
self._device = torch.device('cuda')
elif device in ['mps', 'MPS']:
logging.info("Device is set to mps")
self._device = torch.device('mps')
else:
logging.info("Device is set to CPU")
self._device = torch.device('cpu')
Expand Down Expand Up @@ -454,8 +464,9 @@ def __init__(
model_params: A dictionary of arguments required to instantiate the model
class.
device: the device on which you wish to run the model. If
``device = GPU`` then a GPU device will be used if it is available.
Otherwise, it will be CPU.
``device = GPU`` then a cuda device will be used if it is available.
If ``device = MPS`` then a mps device will be used if it is available.
Otherwise, it will be cpu.
inference_fn: the function to invoke on run_inference.
default = default_keyed_tensor_inference_fn
torch_script_model_path: Path to the torch script model.
Expand All @@ -481,9 +492,12 @@ def __init__(
on torch>=1.9.0,<1.14.0.
"""
self._state_dict_path = state_dict_path
if device == 'GPU':
if device in ['gpu', 'GPU', 'cuda', 'CUDA']:
logging.info("Device is set to CUDA")
self._device = torch.device('cuda')
elif device in ['mps', 'MPS']:
logging.info("Device is set to mps")
self._device = torch.device('mps')
else:
logging.info("Device is set to CPU")
self._device = torch.device('cpu')
Expand Down
Loading