在進行深度學習訓練模型時，對于計算量小一些的模型，是可以在CPU上進行的。但是當計算量比較大時，我們希望利用GPU并行計算的能力去加快訓練的速度。

查看GPU版本號

import torchprint(torch.__version__) ?# 查看torch當前版本號print(torch.version.cuda) ?# 編譯當前版本的torch使用的cuda版本號print(torch.cuda.is_available()) ?# 查看當前cuda是否可用于當前版本的Torch，如果輸出True，則表示可用

查看GPU數量

def try_gpu(i=0): """如果存在，則返回gpu(i)，否則返回cpu()"""if torch.cuda.device_count() >= i + 1:return torch.device(f'cuda:{i}')return torch.device('cpu')def try_all_gpus(): """返回所有可用的GPU，如果沒有GPU，則返回[cpu(),]"""devices = [torch.device(f'cuda:{i}')for i in range(torch.cuda.device_count())]return devices if devices else [torch.device('cpu')] # 0號GPU是否存在，10號GPU是否存在 try_gpu(), try_gpu(10), try_all_gpus()

指定GPU

import torch from torch import nntorch.device('gpu'), torch.cuda.device('cuda'), torch.cuda.device('cuda:1')

GPU計算張量

# 創建一個張量Y在1號GPU Y = torch.rand(2, 3, device=try_gpu(1)) Z = X.cuda(1) # 將X的內容復制在1號GPU的Z print(X) print(Z) tensor([[1., 1., 1.],[1., 1., 1.]], device='cuda:0') tensor([[1., 1., 1.],[1., 1., 1.]], device='cuda:1')

指定GPU計算神經網絡模型

from torch import nnnet = nn.Linear(3, 1) print(list(net.parameters())[0].device) # cpunet.cuda(0) print(list(net.parameters())[0].device) # cuda:0net2 = nn.Linear(3, 1, device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')) print(list(net2.parameters())[0].device) # cuda:0net3 = nn.Linear(3, 1).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu')) print(list(net3.parameters())[0].device) # cuda:0net3 = net3.cpu() print(list(net3.parameters())[0].device) # cpu

總結

以上是生活随笔為你收集整理的pytorch采用GPU加速方法的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。