1. 程式人生 > >淺談將Pytorch模型從CPU轉換成GPU

淺談將Pytorch模型從CPU轉換成GPU

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

環境:Ubuntu 16.04.3
Python版本:3.5.2
Pytorch版本:0.4.0

0. 序言

大家知道,在深度學習中使用GPU來對模型進行訓練是可以通過並行化其計算來提高執行效率,這裡就不多談了。
最近申請到了實驗室的伺服器來跑程式,成功將我簡陋的程式改成了“高大上”GPU版本。
看到網上總體來說少了很多介紹,這裡決定將我的一些思考和工作記錄下來。

1. 如何進行遷移

由於我使用的是Pytorch寫的模型,網上給出了一個非常簡單的轉換方式: 對模型和相應的資料進行.cuda()處理。通過這種方式,我們就可以將記憶體中的資料複製到GPU的視訊記憶體中去

。從而可以通過GPU來進行運算了。

網上說的非常簡單,但是實際使用過程中還是遇到了一些疑惑。下面分資料模型兩方面的遷移來進行說明介紹。

1.1 判定使用GPU

下載了對應的GPU版本的Pytorch之後,要確保GPU是可以進行使用的,通過torch.cuda.is_available()的返回值來進行判斷。返回True則具有能夠使用的GPU。
通過torch.cuda.device_count()可以獲得能夠使用的GPU數量。其他就不多贅述了。
常常通過如下判定來寫可以跑在GPU和CPU上的通用模型:

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

2. 對應資料的遷移

資料方面常用的主要是兩種 —— TensorVariable。實際上這兩種型別是同一個東西,因為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)]

其資料型別會由torch.FloatTensor變為torch.cuda.FloatTensor (GPU 0)這樣代表這個資料現在儲存在
GPU 0的視訊記憶體中了。
如果要將視訊記憶體中的資料複製到記憶體中,則對cuda資料型別使用.cpu()方法即可。

2.2 將Variable遷移到視訊記憶體中去

在模型中,我們最常使用的是Variable這個容器來裝載使用資料。主要是由於Variable可以進行反向傳播來進行自動求導。
同樣地,要將Variable遷移到視訊記憶體中,同樣只需要使用.cuda()即可實現。

這裡有一個小疑問,對Variable直接使用.cuda和對Tensor進行.cuda然後再放置到Variable中結果是否一致呢。答案是肯定的。

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)]

最終我們能發現他們都能夠達到相同的目的,但是他們完全一樣了嗎?我們使用V1 is V2發現,結果是否定的。

對於V1,我們是直接對Variable進行操作的,這樣子V1的.grad_fn中會記錄下建立的方式。因此這二者並不是完全相同的。

2.3 資料遷移小結

.cuda()操作預設使用GPU 0也就是第一張顯示卡來進行操作。當我們想要儲存在其他顯示卡中時可以使用.cuda(<顯示卡號數>)來將資料儲存在指定的顯示卡中。還有很多種方式,具體參考官方文件。

對於不同儲存位置的變數,我們是不可以對他們直接進行計算的。儲存在不同位置中的資料是不可以直接進行互動計算的。
換句話說也就是上面例子中的torch.FloatTensor是不可以直接與torch.cuda.FloatTensor進行基本運算的。位於不同GPU視訊記憶體上的資料也是不能直接進行計算的。

對於Variable,其實就僅僅是一種能夠記錄操作資訊並且能夠自動求導的容器,實際上的關鍵資訊並不在Variable本身,而更應該側重於Variable中儲存的data。

3. 模型遷移

模型的遷移這裡指的是torch.nn下面的一些網路模型以及自己建立的模型遷移到GPU上去。

上面講了使用.cuda()即可將資料從記憶體中移植到視訊記憶體中去。
對於模型來說,也是同樣的方式,我們使用.cuda來將網路放到視訊記憶體上去。

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

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

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

# 建立一個輸入維度為2,輸出維度為2的單層神經網路
linear = torch.nn.Linear(2, 2)
>>>> Linear(in_features=2, out_features=2)

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

我們很驚奇地發現對於模型來說,不像資料那樣使用了.cuda()之後會改變其的資料型別。模型看起來沒有任何的變化。
但是他真的沒有改變嗎。
我們將data1投入linear_cuda中去可以發現,系統會報錯,而將.cuda之後的data2投入linear_cuda才能正常工作。並且輸出的也是具有cuda的資料型別。

那是怎麼一回事呢?
這是因為這些所謂的模型,其實也就是對輸入引數做了一些基本的矩陣運算。所以我們對模型.cuda()實際上也相當於將模型使用到的引數儲存到了視訊記憶體上去。

對於上面的例子,我們可以通過觀察引數來發現區別所在。

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 自己模型的遷移

對於自己建立的模型類,由於繼承了torch.nn.Module,則可同樣使用.cuda()來將模型中用到的所有引數都儲存到視訊記憶體中去。

這裡筆者曾經有一個疑問:當我們對模型儲存到視訊記憶體中去之後,那麼這個模型中的方法後面所創建出來的Tensor是不是都會預設變成cuda的資料型別。答案是否定的。具體操作留給讀者自己去實現。

3.3 模型小結

對於模型而言,我們可以將其看做是一種類似於Variable的容器。我們對它進行.cuda()處理,是將其中的引數放到視訊記憶體上去(因為實際使用的時候也是通過這些引數做運算)。

4. 總結

Pytorch使用起來直接簡單,GPU的使用也是簡單明瞭。然而對於多GPU和CPU的協同使用則還是有待提高。