测试代码
import torch
import time
a = torch.rand((10000, 10000))
b = torch.rand((10000, 10000))
start_time_cpu = time.time()
c = a.matmul(b)
end_time_cpu = time.time()
dtime_cpu = end_time_cpu - start_time_cpu
a = a.cuda()
b = b.cuda()
start_time_gpu = time.time()
c = a.matmul(b)
end_time_gpu = time.time()
dtime_gpu = end_time_gpu - start_time_gpu
ratio = dtime_cpu / dtime_gpu
print(f"cpu: {dtime_cpu}秒\ngpu: {dtime_gpu}秒\nradio: {ratio}")
执行结果
cpu: 7.425754547119141秒
gpu: 0.08930134773254395秒
radio: 83.15389113005497