注意Pytorch中GPU與CPU的使用

其實在Pytorch中並不是在任何情況下把數據放到GPU下就能訓練效率，減少訓練時間。

下面舉個例子：

import torch
import time

print(torch.cuda.is_available())

print(torch.cuda.current_device())



###CPU
start_time = time.time()
a = torch.ones(5,5)
for _ in range(1000000):
    a += a
elapsed_time = time.time() - start_time

print('CPU time = ',elapsed_time)

###GPU
start_time = time.time()
b = torch.ones(5,5).cuda()
for _ in range(1000000):
    b += b
elapsed_time = time.time() - start_time

print('GPU time = ',elapsed_time)

運行結果：

True
0
CPU time =  1.8500258922576904
GPU time =  8.387852907180786

Process finished with exit code 0

可以看出CPU運行的速度比GPU快很多，這是什麼原因呢?

其實，這是由於GPU加速通過大量的計算並行化來工作.即使同步計算，然而在上面的例子同步計算只有矩陣中25個元素，對於for循環而言，必須知道前一個時候值才能進行下一時刻值的計算，所以並不能實現並行計算。

但是如果改變張量的大小,事情會發生變化,那麼PyTorch能夠並行化更多的整體計算.我將迭代更改爲500的值,結果如下：

True
0
CPU time =  22.1019504070282
GPU time =  86.17750453948975

Process finished with exit code 0

可以看出當爲當爲5×5的時候CPU爲GPU的大約4.5倍，當爲500×500時只有3.89倍，這裏可以更改爲更大的值你就會發現GPU的效果了，所以以後知道什麼時候用CPU,什麼時候用GPU了吧。