MPS 後端¶
mps 裝置可在搭載 Metal 程式設計框架的 MacOS 裝置上進行高效能訓練。它引入了一個新的裝置,將機器學習計算圖和原語分別映射到高效的 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)