MPS 后端

MPS 后端

伊织code2024-07-17 18:36

本文来自：
https://pytorch.org/docs/stable/notes/mps.html
https://pytorch.ac.cn/docs/stable/notes/mps.html

mps 设备支持在使用 Metal 编程框架的 MacOS 设备上，进行高性能 GPU 训练。

它引入了新的设备，将机器学习计算图和原语映射到 Metal Performance Shaders 图框架和 Metal Performance Shaders 框架提供的经过优化的内核上。

新的 MPS 后端扩展了 PyTorch 生态系统，并为现有脚本提供在 GPU 上设置和运行操作的功能。

要开始使用，只需将您的张量和模块移动到 mps 设备。

python 复制代码

# 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)

2024-07-16（二）