Shortcuts

MPS 后端

mps 设备使 MacOS 设备上的 GPU 能够进行高性能训练,使用的是 Metal 编程框架。它引入了一个新设备,用于在高效的 Metal Performance Shaders Graph 框架和由 Metal Performance Shaders 框架提供的调优内核上分别映射机器学习计算图和原语。

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

要开始使用,只需将您的 Tensor 和 Module 移动到 mps 设备:

# 检查MPS是否可用
if not torch.backends.mps.is_available():
    if not torch.backends.mps.is_built():
        print("MPS不可用,因为当前的PyTorch安装未启用MPS。")
    else:
        print("MPS不可用,因为当前的MacOS版本不是12.3+,和/或您没有此机器上的MPS启用设备。")

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

    # 直接在mps设备上创建一个Tensor
    x = torch.ones(5, device=mps_device)
    # 或者
    x = torch.ones(5, device="mps")

    # 任何操作都在GPU上进行
    y = x * 2

    # 将您的模型移动到mps设备,就像其他设备一样
    model = YourFavoriteNet()
    model.to(mps_device)

    # 现在每次调用都在GPU上运行
    pred = model(x)