【深度学习】如何在PyToch中查看是否存在mps(针对M系列芯片)

最新推荐文章于 2024-12-21 14:35:57 发布

Lydia.na

最新推荐文章于 2024-12-21 14:35:57 发布

阅读量696

点赞数

文章标签：深度学习人工智能

原文链接：https://pytorch.org/docs/stable/notes/mps.html

版权

【深度学习】如何在PyToch中查看是否存在mps(针对M系列芯片)

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)

我根据我所训练的模型，将其更改：

def try_gpu():  #@save
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    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.")
            mps_device = torch.device("cpu")
        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.")
            mps_device = torch.device("cpu")

    else:
        mps_device = torch.device("mps")