如何在Mac上使用Metal加速PyTorch

概述

之前在PC上,我们可以使用CUDA进行AI训练加速,但是在Mac上却只能使用CPU。现在苹果联合PyTorch,推出了Metal作为PyTorch的计算后端,苹果的文档描述如下

Metal backend for PyTorch
The new Metal backend in PyTorch version 1.12 enables high-performance, GPU-accelerated training using MPS Graph and the Metal Performance Shaders primitives.

我们只需要在Mac上安装1.12版本以上的PyTorch,就可以使用MPS作为后端,享受Metal带来的计算加速了

使用conda安装PyTorch

基于conda安装torch和图像音频处理库

conda install pytorch torchvision torchaudio -c pytorch

测试MPS是否可用

编写一个简单的例子测试MPS是否可用,在mps设备上创建一个tensor变量,然后打印出来

import torch

device = torch.device("mps")

val = torch.rand((2,2), device=device)
print(val)

如果一切没问题,会打印出来类似的结果

tensor([[0.9615, 0.5488],
        [0.1987, 0.4467]], device='mps:0')
©著作权归作者所有,转载或内容合作请联系作者
【社区内容提示】社区部分内容疑似由AI辅助生成,浏览时请结合常识与多方信息审慎甄别。
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

相关阅读更多精彩内容

友情链接更多精彩内容