快捷導(dǎo)航

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

更新時(shí)間：2022年07月15日 15:10:21 作者：康康同學(xué)97

這篇文章主要介紹了pytorch中Tensor.to(device)和model.to(device)的區(qū)別及說(shuō)明，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。如有錯(cuò)誤或未考慮完全的地方，望不吝賜教

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

區(qū)別所在

使用GPU訓(xùn)練的時(shí)候，需要將Module對(duì)象和Tensor類(lèi)型的數(shù)據(jù)送入到device。通常會(huì)使用 to.(device)。但是需要注意的是：

對(duì)于Tensor類(lèi)型的數(shù)據(jù)，使用to.(device) 之后，需要接收返回值，返回值才是正確設(shè)置了device的Tensor。
對(duì)于Module對(duì)象，只用調(diào)用to.(device) 就可以將模型設(shè)置為指定的device。不必接收返回值。

來(lái)自pytorch官方文檔的說(shuō)明：

Tensor.to(device)

Module.to(device)

舉例

# Module對(duì)象設(shè)置device的寫(xiě)法
model.to(device)

# Tensor類(lèi)型的數(shù)據(jù)設(shè)置 device 的寫(xiě)法。
samples = samples.to(device)

pytorch學(xué)習(xí)筆記--to(device)用法

在學(xué)習(xí)深度學(xué)習(xí)的時(shí)候，我們寫(xiě)代碼經(jīng)常會(huì)見(jiàn)到類(lèi)似的代碼：

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)

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

這段代碼的意思就是將所有最開(kāi)始讀取數(shù)據(jù)時(shí)的tensor變量copy一份到device所指定的GPU上去，之后的運(yùn)算都在GPU上進(jìn)行。

為什么要在GPU上做運(yùn)算呢？

首先，在做高維特征運(yùn)算的時(shí)候，采用GPU無(wú)疑是比用CPU效率更高，如果兩個(gè)數(shù)據(jù)中一個(gè)加了.cuda()或者.to(device)，而另外一個(gè)沒(méi)有加，就會(huì)造成類(lèi)型不匹配而報(bào)錯(cuò)。

tensor和numpy都是矩陣，前者能在GPU上運(yùn)行，后者只能在CPU運(yùn)行，所以要注意數(shù)據(jù)類(lèi)型的轉(zhuǎn)換。

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

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

如果你有多個(gè)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)

以上為個(gè)人經(jīng)驗(yàn)，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: