分别使用CPU和GPU进行Pytorch中的Tensor(张量)计算,测试Tensor在两种不同运算设备上的计算速度差异。

设备:

服务器:Dell EMC Power Edge R740

CPU:Intel Xeon Gold 5117 * 2

Memory:64G

GPU:NVIDIA Tesla T4 16G * 1

Python Version:3.8

CUDA Version:11.4

Pytorch Version:1.9.0

分别使用三种不同尺寸的Tensor进行平方运算测试,单次测试进行10万次平方运算,分别在CPU和GPU上进行20次测试,时长取20次测试的平均结果。

Code:



import warnings
warnings.filterwarnings('ignore')


import torch
import time


if torch.cuda.is_available():
    device=torch.device('cuda:0')
    print('The current device is GPU. ',end='\n\n')
else:
    device=torch.device('cpu')
    print('The current device is CPU. ',end='\n\n')


a=torch.normal(mean=0, std=1, size=(32,128,128))
    

b=a.clone()
d=a.clone()
d=d.to(device)


Test_times=20


time_cost=0
for _ in range(Test_times):
    time_0=time.time()
    for i in range(100000):
        c=b**2
    time_1=time.time()
    
    time_cost=time_cost + time_1-time_0
time_cost=time_cost/Test_times
    
print(f'Average CPU Time : {time_cost:.5f} ')


time_cost=0
for _ in range(Test_times):
    time_2=time.time()
    for i in range(100000):
        e=d**2
    time_3=time.time()
    
    time_cost=time_cost + time_3-time_2
time_cost=time_cost/Test_times

print(f'Average GPU Time : {time_cost:.5f} ')





Result:

1.    32*64*64

2.    32*128*128

3.    32*256*256

Analysis:

从测试结果来看,无论何种尺寸的张量计算,在GPU上的运算速度都要远快于CPU。

 

 

 

Logo

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

更多推荐