pytorch中Tensor.to(device)和model.to(device)的區別及說明

Posted on by

Tensor.to(device)和model.to(device)的區別

區別所在

使用GPU訓練的時候，需要將Module對象和Tensor類型的數據送入到device。通常會使用 to.(device)。但是需要註意的是：

  • 對於Tensor類型的數據，使用to.(device) 之後，需要接收返回值，返回值才是正確設置瞭device的Tensor。
  • 對於Module對象，隻用調用to.(device) 就可以將模型設置為指定的device。不必接收返回值。

來自pytorch官方文檔的說明：

Tensor.to(device)

Module.to(device)

舉例

# Module對象設置device的寫法
model.to(device)

# Tensor類型的數據設置 device 的寫法。
samples = samples.to(device)

pytorch學習筆記–to(device)用法

在學習深度學習的時候，我們寫代碼經常會見到類似的代碼：

img = img.to(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))
model = models.vgg16_bn(pretrained=True).to(device=torch.device("cuda" if torch.cuda.is_available() else "cpu"))

也可以先定義device：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
img = img.to(device)

這段代碼到底有什麼用呢？

這段代碼的意思就是將所有最開始讀取數據時的tensor變量copy一份到device所指定的GPU上去，之後的運算都在GPU上進行。

為什麼要在GPU上做運算呢？

首先，在做高維特征運算的時候，采用GPU無疑是比用CPU效率更高，如果兩個數據中一個加瞭.cuda()或者.to(device)，而另外一個沒有加，就會造成類型不匹配而報錯。

tensor和numpy都是矩陣，前者能在GPU上運行，後者隻能在CPU運行，所以要註意數據類型的轉換。

.cuda()和.to(device)的效果一樣嗎？為什麼後者更好？

兩個方法都可以達到同樣的效果，在pytorch中，即使是有GPU的機器，它也不會自動使用GPU，而是需要在程序中顯示指定。調用model.cuda()，可以將模型加載到GPU上去。這種方法不被提倡，而建議使用model.to(device)的方式，這樣可以顯示指定需要使用的計算資源，特別是有多個GPU的情況下。

如果你有多個GPU

那麼可以參考以下代碼：

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

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model，device_ids=[0,1,2])
    model.to(device)

以上為個人經驗，希望能給大傢一個參考，也希望大傢多多支持WalkonNet。

推薦閱讀：