淺談將Pytorch模型從CPU轉換成GPU
最近將Pytorch程式遷移到GPU上去的一些工作和思考
環境:Ubuntu 16.04.3
Python版本:3.5.2
Pytorch版本:0.4.0
0. 序言
大家知道,在深度學習中使用GPU來對模型進行訓練是可以通過並行化其計算來提高執行效率,這裡就不多談了。
最近申請到了實驗室的伺服器來跑程式,成功將我簡陋的程式改成了“高大上”GPU版本。
看到網上總體來說少了很多介紹,這裡決定將我的一些思考和工作記錄下來。
1. 如何進行遷移
由於我使用的是Pytorch寫的模型,網上給出了一個非常簡單的轉換方式: 對模型和相應的資料進行.cuda()處理。通過這種方式,我們就可以將記憶體中的資料複製到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. 對應資料的遷移
資料方面常用的主要是兩種 —— 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)]
其資料型別會由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的協同使用則還是有待提高。