將Pytorch模型從CPU轉(zhuǎn)換成GPU的實現(xiàn)方法

最近將Pytorch程序遷移到GPU上去的一些工作和思考

環(huán)境：Ubuntu 16.04.3

Python版本：3.5.2

Pytorch版本：0.4.0

0. 序言

大家知道，在深度學習中使用GPU來對模型進行訓練是可以通過并行化其計算來提高運行效率，這里就不多談了。

最近申請到了實驗室的服務器來跑程序，成功將我簡陋的程序改成了“高大上”GPU版本。

看到網(wǎng)上總體來說少了很多介紹，這里決定將我的一些思考和工作記錄下來。

1. 如何進行遷移

由于我使用的是Pytorch寫的模型，網(wǎng)上給出了一個非常簡單的轉(zhuǎn)換方式: 對模型和相應的數(shù)據(jù)進行.cuda()處理。通過這種方式，我們就可以將內(nèi)存中的數(shù)據(jù)復制到GPU的顯存中去。從而可以通過GPU來進行運算了。

網(wǎng)上說的非常簡單，但是實際使用過程中還是遇到了一些疑惑。下面分數(shù)據(jù)和模型兩方面的遷移來進行說明介紹。

1.1 判定使用GPU

下載了對應的GPU版本的Pytorch之后，要確保GPU是可以進行使用的，通過torch.cuda.is_available()的返回值來進行判斷。返回True則具有能夠使用的GPU。

通過torch.cuda.device_count()可以獲得能夠使用的GPU數(shù)量。其他就不多贅述了。

常常通過如下判定來寫可以跑在GPU和CPU上的通用模型：

if torch.cuda.is_available():
  ten1 = ten1.cuda()
  MyModel = MyModel.cuda() 

2. 對應數(shù)據(jù)的遷移

數(shù)據(jù)方面常用的主要是兩種 —— Tensor和Variable。實際上這兩種類型是同一個東西，因為Variable實際上只是一個容器，這里先視其不同。

2.1 將Tensor遷移到顯存中去

不論是什么類型的Tensor（FloatTensor或者是LongTensor等等），一律直接使用方法.cuda()即可。

例如：

ten1 = torch.FloatTensor(2)
>>>> 6.1101e+24
   4.5659e-41
   [torch.FloatTensor of size 2]

ten1_cuda = ten1.cuda()
>>>>  6.1101e+24
    4.5659e-41
    [torch.cuda.FloatTensor of size 2 (GPU 0)]

其數(shù)據(jù)類型會由torch.FloatTensor變?yōu)閠orch.cuda.FloatTensor （GPU 0）這樣代表這個數(shù)據(jù)現(xiàn)在存儲在

GPU 0的顯存中了。

如果要將顯存中的數(shù)據(jù)復制到內(nèi)存中，則對cuda數(shù)據(jù)類型使用.cpu()方法即可。

2.2 將Variable遷移到顯存中去

在模型中，我們最常使用的是Variable這個容器來裝載使用數(shù)據(jù)。主要是由于Variable可以進行反向傳播來進行自動求導。

同樣地，要將Variable遷移到顯存中，同樣只需要使用.cuda()即可實現(xiàn)。

這里有一個小疑問，對Variable直接使用.cuda和對Tensor進行.cuda然后再放置到Variable中結(jié)果是否一致呢。答案是肯定的。

ten1 = torch.FloatTensor(2)
>>> 6.1101e+24
   4.5659e-41
  [torch.FloatTensor of size 2]

ten1_cuda = ten1.cuda()
>>>> 6.1101e+24
   4.5659e-41
  [torch.cuda.FloatTensor of size 2 (GPU 0)]

V1_cpu = autograd.Variable(ten1)
>>>> Variable containing:
   6.1101e+24
   4.5659e-41
  [torch.FloatTensor of size 2]

V2 = autograd.Variable(ten1_cuda)
>>>> Variable containing:
   6.1101e+24
   4.5659e-41
  [torch.cuda.FloatTensor of size 2 (GPU 0)]

V1 = V1_cpu.cuda()
>>>> Variable containing:
   6.1101e+24
   4.5659e-41
  [torch.cuda.FloatTensor of size 2 (GPU 0)]

最終我們能發(fā)現(xiàn)他們都能夠達到相同的目的，但是他們完全一樣了嗎？我們使用V1 is V2發(fā)現(xiàn)，結(jié)果是否定的。

對于V1，我們是直接對Variable進行操作的,這樣子V1的.grad_fn中會記錄下創(chuàng)建的方式。因此這二者并不是完全相同的。

2.3 數(shù)據(jù)遷移小結(jié)

.cuda()操作默認使用GPU 0也就是第一張顯卡來進行操作。當我們想要存儲在其他顯卡中時可以使用.cuda(<顯卡號數(shù)>)來將數(shù)據(jù)存儲在指定的顯卡中。還有很多種方式，具體參考官方文檔。

對于不同存儲位置的變量，我們是不可以對他們直接進行計算的。存儲在不同位置中的數(shù)據(jù)是不可以直接進行交互計算的。

換句話說也就是上面例子中的torch.FloatTensor是不可以直接與torch.cuda.FloatTensor進行基本運算的。位于不同GPU顯存上的數(shù)據(jù)也是不能直接進行計算的。

對于Variable，其實就僅僅是一種能夠記錄操作信息并且能夠自動求導的容器，實際上的關鍵信息并不在Variable本身，而更應該側(cè)重于Variable中存儲的data。

3. 模型遷移

模型的遷移這里指的是torch.nn下面的一些網(wǎng)絡模型以及自己創(chuàng)建的模型遷移到GPU上去。

上面講了使用.cuda()即可將數(shù)據(jù)從內(nèi)存中移植到顯存中去。

對于模型來說，也是同樣的方式，我們使用.cuda來將網(wǎng)絡放到顯存上去。

3.1 torch.nn下的基本模型遷移

這里使用基本的單層感知機來進行舉例(線性模型)。

data1 = torch.FloatTensor(2)
data2 = data1.cuda

# 創(chuàng)建一個輸入維度為2，輸出維度為2的單層神經(jīng)網(wǎng)絡
linear = torch.nn.Linear(2, 2)
>>>> Linear(in_features=2, out_features=2)

linear_cuda = linear.cuda()
>>>> Linear(in_features=2, out_features=2)

我們很驚奇地發(fā)現(xiàn)對于模型來說，不像數(shù)據(jù)那樣使用了.cuda()之后會改變其的數(shù)據(jù)類型。模型看起來沒有任何的變化。

但是他真的沒有改變嗎。

我們將data1投入linear_cuda中去可以發(fā)現(xiàn)，系統(tǒng)會報錯，而將.cuda之后的data2投入linear_cuda才能正常工作。并且輸出的也是具有cuda的數(shù)據(jù)類型。

那是怎么一回事呢？

這是因為這些所謂的模型，其實也就是對輸入?yún)?shù)做了一些基本的矩陣運算。所以我們對模型.cuda()實際上也相當于將模型使用到的參數(shù)存儲到了顯存上去。

對于上面的例子，我們可以通過觀察參數(shù)來發(fā)現(xiàn)區(qū)別所在。

linear.weight
>>>> Parameter containing:
  -0.6847 0.2149
  -0.5473 0.6863
  [torch.FloatTensor of size 2x2]

linear_cuda.weight
>>>> Parameter containing:
  -0.6847 0.2149
  -0.5473 0.6863
  [torch.cuda.FloatTensor of size 2x2 (GPU 0)]

3.2 自己模型的遷移

對于自己創(chuàng)建的模型類，由于繼承了torch.nn.Module，則可同樣使用.cuda()來將模型中用到的所有參數(shù)都存儲到顯存中去。

這里筆者曾經(jīng)有一個疑問：當我們對模型存儲到顯存中去之后，那么這個模型中的方法后面所創(chuàng)建出來的Tensor是不是都會默認變成cuda的數(shù)據(jù)類型。答案是否定的。具體操作留給讀者自己去實現(xiàn)。

3.3 模型小結(jié)

對于模型而言，我們可以將其看做是一種類似于Variable的容器。我們對它進行.cuda()處理，是將其中的參數(shù)放到顯存上去（因為實際使用的時候也是通過這些參數(shù)做運算）。

4. 總結(jié)

Pytorch使用起來直接簡單，GPU的使用也是簡單明了。然而對于多GPU和CPU的協(xié)同使用則還是有待提高。

以上這篇將Pytorch模型從CPU轉(zhuǎn)換成GPU的實現(xiàn)方法就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持億速云。