pytorch使用gpu加速的方法

一、默認gpu加速

一般來說我們最常見到的用法是這樣的：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

或者說:

if torch.cuda.is_available():
	device = torch.device("cuda")
else:
	device = torch.device("cpu")

這樣我們就可以把某一個向量或者模型進行gpu訓練

data = data.to(device)
model = Model(...).to(device)

二、指定gpu加速

torch.device('cuda', 0)
torch.device('cuda:0')

來指定使用的具體設備。如果沒有顯式指定設備序號的話則使用torch.cuda.current_device()對應的序號。