概述
之前在PC上,我們可以使用CUDA進(jìn)行AI訓(xùn)練加速,但是在Mac上卻只能使用CPU。現(xiàn)在蘋果聯(lián)合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
設(shè)備上創(chuàng)建一個tensor變量,然后打印出來
import torch
device = torch.device("mps")
val = torch.rand((2,2), device=device)
print(val)
如果一切沒問題,會打印出來類似的結(jié)果
tensor([[0.9615, 0.5488],
[0.1987, 0.4467]], device='mps:0')