MPS backend#
Created On: May 13, 2022 | Last Updated On: Jun 02, 2022
mps device enables high-performancetraining on GPU for MacOS devices with Metal programming framework. Itintroduces a new device to map Machine Learning computational graphs andprimitives on highly efficient Metal Performance Shaders Graph framework andtuned kernels provided by Metal Performance Shaders framework respectively.
The new MPS backend extends the PyTorch ecosystem and provides existing scriptscapabilities to setup and run operations on GPU.
To get started, simply move your Tensor and Module to themps device:
# Check that MPS is availableifnottorch.backends.mps.is_available():ifnottorch.backends.mps.is_built():print("MPS not available because the current PyTorch install was not ""built with MPS enabled.")else:print("MPS not available because the current MacOS version is not 12.3+ ""and/or you do not have an MPS-enabled device on this machine.")else:mps_device=torch.device("mps")# Create a Tensor directly on the mps devicex=torch.ones(5,device=mps_device)# Orx=torch.ones(5,device="mps")# Any operation happens on the GPUy=x*2# Move your model to mps just like any other devicemodel=YourFavoriteNet()model.to(mps_device)# Now every call runs on the GPUpred=model(x)