在 PyTorch 中以下數(shù)據(jù)結(jié)構(gòu)分為 CPU 和 GPU 兩個版本：

• Tensor
• nn.Module （包括常用的 layer 、loss function ，以及容器 Sequential 等）

它們都帶有一個 .cuda 方法，調(diào)用此方法即可將其轉(zhuǎn)為對應的 GPU 對象。

注意，tensor.cuda 會返回一個新對象，這個新對象的數(shù)據(jù)已轉(zhuǎn)移至GPU，而之前的 tensor 還在原來的設備上（CPU ）。而 module.cuda 則會將所有的數(shù)據(jù)都遷移至 GPU ，并返回自己。所以 module = module.cuda() 和 module.cuda() 所起的作用一致。

tensor = t.Tensor(3, 4)
# 返回一個新的tensor，但原來的tensor并沒有改變
tensor.cuda(0)
tensor.is_cuda  # False

重新賦給自己，tensor 指向 cuda 上的數(shù)據(jù)，不再執(zhí)行原數(shù)據(jù)。不指定使用的 GPU 設備，將默認使用第 1 塊 GPU。

tensor = tensor.cuda()
tensor.is_cuda	# True

nn.Module 在 GPU 與 CPU 之間的轉(zhuǎn)換，本質(zhì)上還是利用了 Tensor 在 GPU 和 CPU 之間的轉(zhuǎn)換。nn.Module 的 cuda 方法是將 nn.Module 下的所有 parameter （包括子 module 的 parameter ）都轉(zhuǎn)移至GPU ，而 Parameter 本質(zhì)上也是 tensor ( Tensor 的子類)。

# 將nn.Module模型放到cuda上，其子模型也都自動放到cuda上
from torch import nn
module = nn.Linear(3, 4)
module.cuda(device = 0)
module.weight.is_cuda	# True

注意： 為什么將數(shù)據(jù)轉(zhuǎn)移至 GPU 的方法叫做 .cuda 而不是 .gpu ，就像將數(shù)據(jù)轉(zhuǎn)移至 CPU 調(diào)用的方法是 .cpu ？
這是因為 GPU 的編程接口采用 CUDA ，而目前并不是所有的 GPU 都支持 CUDA ，只有部分 Nvidia 的GPU 才支持。PyTorch 未來可能會支持 AMD 的 GPU ，而 AMD GPU 的編程接口采用 OpenCL ，因此PyTorch 還預留著 .cl 方法，用于以后支持 AMD 等的 GPU 。

下面將舉例說明，這部分代碼需要你具有兩塊GPU設備。

class VeryBigModule(nn.Module):def __init__(self):super(VeryBigModule, self).__init__()self.GiantParameter1 = t.nn.Parameter(t.randn(100000, 20000)).cuda(0)self.GiantParameter2 = t.nn.Parameter(t.randn(20000, 100000)).cuda(1)def forward(self, x):x = self.GiantParameter1.mm(x.cuda(0))x = self.GiantParameter2.mm(x.cuda(1))return x

上面最后一部分中，兩個 Parameter 所占用的內(nèi)存空間都非常大，大概是 8 個 G，如果將這兩個都同時放在一塊 GPU 上幾乎會將顯存占滿，無法再進行任何其它運算。此時可通過這種方式將不同的計算分布到不同的 GPU 中。

關(guān)于使用 GPU 的一些建議：

• GPU 運算很快，但對于很小的運算量來說，并不能體現(xiàn)出它的優(yōu)勢，因此對于一些簡單的操作可直接利用CPU 完成；
• 數(shù)據(jù)在 CPU 和 GPU 之間，以及 GPU 與 GPU 之間的傳遞會比較耗時，應當盡量避免；
• 在進行低精度的計算時，可以考慮 HalfTensor ，它相比于 FloatTensor 能節(jié)省一半的顯存，但需千萬注意數(shù)值溢出的情況；

而除了調(diào)用對象的 .cuda 方法之外，還可以使用 torch.cuda.device ，來指定默認使用哪一塊 GPU ，或使用torch.set_default_tensor_type 使程序默認使用 GPU ，不需要手動調(diào)用 cuda 。

# 如果未指定使用哪塊GPU，默認使用GPU 0
x = t.cuda.FloatTensor(2, 3)
# x.get_device() == 0
y = t.FloatTensor(2, 3).cuda()
# y.get_device() == 0# 指定默認使用GPU 0
with t.cuda.device(0):    # 在GPU 0上構(gòu)建tensora = t.cuda.FloatTensor(2, 3)# 將tensor轉(zhuǎn)移至GPU 0b = t.FloatTensor(2, 3).cuda()print(a.get_device() == b.get_device() == 0 )	# Truec = a + bprint(c.get_device() == 0)	# Truez = x + yprint(z.get_device() == 0)	# True# 手動指定使用GPU 0d = t.randn(2, 3).cuda(0)print(d.get_device() == 2)	# False

# 指定默認tensor的類型為GPU上的FloatTensor
t.set_default_tensor_type('torch.cuda.FloatTensor') 
a = t.ones(2, 3)
a.is_cuda # True

如果服務器具有多個 GPU ，tensor.cuda() 方法會將 tensor 保存到第一塊 GPU 上，等價于tensor.cuda(0) 。此時如果想使用第二塊 GPU ，需手動指定 tensor.cuda(1) ，而這需要修改大量代碼，很是繁瑣。這里有兩種替代方法：

• 一種是先調(diào)用 t.cuda.set_device(1) 指定使用第二塊 GPU ，后續(xù)的 .cuda() 都無需更改，切換 GPU 只需修改這一行代碼。
• 更推薦的方法是設置環(huán)境變量 CUDA_VISIBLE_DEVICES ，例如當 export CUDA_VISIBLE_DEVICE=1 （下標是從 0 開始，1 代表第二塊 GPU ），只使用第二塊物理 GPU ，但在程序中這塊 GPU 會被看成是第一塊邏輯 GPU ，因此此時調(diào)用 tensor.cuda() 會將 Tensor 轉(zhuǎn)移至第二塊物理 GPU。CUDA_VISIBLE_DEVICES 還可以指定多個 GPU ，如 export CUDA_VISIBLE_DEVICES=0,2,3 ，那么第一、三、四塊物理 GPU 會被映射成第一、二、三塊邏輯 GPU ，tensor.cuda(1) 會將 Tensor 轉(zhuǎn)移到第三塊物理 GPU 上。

設置 CUDA_VISIBLE_DEVICES 有兩種方法：

• 一種是在命令行中 CUDA_VISIBLE_DEVICES=0,1 python main.py ；
• 一種是在程序中

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

如果使用 IPython 或者 Jupyter notebook，還可以使用 %env CUDA_VISIBLE_DEVICES=1,2 來設置環(huán)境變量。

從 0.4 版本開始，PyTorch 新增了 tensor.to(device) 方法，能夠?qū)崿F(xiàn)設備透明，便于實現(xiàn) CPU/GPU 兼容。

x = torch.randn(1)
if torch.cuda.is_available():device = torch.device("cuda")            y = torch.ones_like(x, device=device)    # directly create a tensor on GPUx = x.to(device)                         z = x + yprint(z)print(z.to("cpu", torch.double))  

輸出：

tensor([ 2.0718], device='cuda:0')
tensor([ 2.0718], dtype=torch.float64)

PyTorch 支持分布式 GPU 。分布式是指有多個 GPU 在多臺服務器上，而并行一般指的是一臺服務器上的多個 GPU。分布式涉及到了服務器之間的通信，因此比較復雜，PyTorch 封裝了相應的接口，可以用幾句簡單的代碼實現(xiàn)分布式訓練。分布式對普通用戶來說比較遙遠，因為搭建一個分布式集群的代價十分大，使用也比較復雜。相比之下一機多卡更加現(xiàn)實。

總結(jié)

以上是生活随笔為你收集整理的PyTorch 笔记（19）— Tensor 用 GPU 加速的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。