參考

4.6 GPU計算

到目前為止,我們一直使用CPU進行計算。對復雜的神經網絡和大規模數據來說,使用CPU來計算可能不夠高效。
在本節中,將要介紹如何使用單塊NIVIDA GPU進行計算

4.6.1 計算設備

PyTorch可以指定用來存儲和計算的設備,如果用內存的CPU或者顯存的GPU。默認情況下,PyTorch會將數據創建在內存,然后利用GPU來計算。
用torch.cuda.is_available()查看GPU是否可用:

import torch from torch import nntorch.cuda.is_available()

# 查看GPU數量 torch.cuda.current_device()

PS: GPU的數序號從0開始計算

# 根據索引查看GPU名字 torch.cuda.get_device_name(0)

4.6.2 Tensor的GPU計算

默認情況下,Tensor會被存在內存上。因此,之前我們每次打印Tensor的時候看不到GPU相關標識

x = torch.tensor([1,2,3]) x

使用.cuda()可以將CPU上的Tensor轉換(復制)到GPU上。如果有多塊GPU,我們用.cuda(i)來表示第i塊GPU

x = x.cuda()

我們可以通過Tensor的device屬性來查看該Tensor所在的設備

x.device

我們可以在創建的時候就指定設備

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')x = torch.tensor([1, 2, 3], device=device) print(x)x= torch.tensor([1,2,3]).to(device)print(x)

如果對在GPU上的數據進行運算，那么結果還是存放在GPU上

y = x**2 y

需要注意的是，存儲在不同位置中的數據是不可以直接進行計算的。即存放在CPU上的數據不可以直接與存放在GPU上的數據進行運算，位于不同GPU上的數據也是不能直接進行計算的

4.6.3 模型的GPU計算

同Tensor類似,PyTorch模型也可以通過 .cuda轉換到GPU上。我們可以通過device屬性來查看存放模型的設備

net = nn.Linear(3, 1) list(net.parameters())[0].device

將模型轉到GPU上

net.cuda() list(net.parameters())[0].device

模型和參數需要同時在GPU上才能計算

x =torch.rand(2, 3).cuda() net(x)

總結

以上是生活随笔為你收集整理的[pytorch、学习] - 4.6 GPU计算的全部內容，希望文章能夠幫你解決所遇到的問題。

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