在跑神经网络的时候,GPU的作用是很明显的。下面比较一下CPU跑和GPU跑的区别:

先用CPU跑:

import torch
import time

for i in range(1,10):
    start_time = time.time() #返回当前语句运行时的时间点(单位秒)。
    a = torch.rand(i*100, 1000, 1000)   #生成i*100个的10001000矩阵组成的张量,每个元素为0-1之间的随机值。
    a = torch.mul(a,a)
    Cul_end_time = time.time() - start_time
    print(Cul_end_time)

输出:

0.441303014755249
0.8848202228546143
1.287822961807251
1.7319905757904053
2.177103281021118
2.555281162261963
2.9773526191711426
3.3875885009765625
3.9365322589874268

Process finished with exit code 0

再用GPU跑,代码如下:

import torch
import time

for i in range(1,10):
    start_time = time.time() #return current time in second.
    a = torch.rand(i*100, 1000, 1000)
    
    a = a.cuda()   #和上面的区别是,a.cuda()把张量a转到GPU中,下面的运算就在GPU进行了
    
    a = torch.mul(a,a)
    Cul_end_time = time.time() - start_time
    print(Cul_end_time)

输出:

526.2110984325409
1.0246520042419434
1.4614639282226562
1.8706181049346924
2.5586912631988525
3.071241617202759
3.5114991664886475
3.9417364597320557
4.653223752975464

Process finished with exit code 0

第一次运算的时候花了526秒,近9分钟,是在等数据加载到显卡中。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐