MPS 後端#
建立時間:2022 年 5 月 13 日 | 最後更新時間:2022 年 6 月 2 日
mps 裝置支援在配備 Metal 程式設計框架的 MacOS 裝置上進行高效能 GPU 訓練。它引入了一個新裝置,該裝置可分別將機器學習計算圖和原語對映到高效的 Metal Performance Shaders Graph 框架和 Metal Performance Shaders 框架提供的最佳化核心。
新的 MPS 後端擴充套件了 PyTorch 生態系統,併為現有指令碼提供了在 GPU 上設定和執行操作的能力。
要開始使用,只需將您的 Tensor 和 Module 移動到 mps 裝置即可
# Check that MPS is available
if not torch.backends.mps.is_available():
if not torch.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 device
x = torch.ones(5, device=mps_device)
# Or
x = torch.ones(5, device="mps")
# Any operation happens on the GPU
y = x * 2
# Move your model to mps just like any other device
model = YourFavoriteNet()
model.to(mps_device)
# Now every call runs on the GPU
pred = model(x)