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)