GPU性能的简单测试脚本（pytorch版）

wonder-wall

已于 2022-10-28 20:30:53 修改

阅读量1.2w

点赞数 8

分类专栏：小技巧文章标签： pytorch 深度学习 python

于 2022-08-30 08:34:18 首次发布

本文链接：https://blog.csdn.net/qq_41129489/article/details/126596108

版权

小技巧专栏收录该内容

8 篇文章 6 订阅

订阅专栏

import time
import torch

# 测试gpu计算耗时
A = torch.ones(5000, 5000).to('cuda')
B = torch.ones(5000, 5000).to('cuda')
startTime2 = time.time()
for i in range(100):
    C = torch.matmul(A, B)
endTime2 = time.time()
print('gpu计算总时长:', round((endTime2 - startTime2) * 1000, 2), 'ms')

# 测试cpu计算耗时
A = torch.ones(5000, 5000)
B = torch.ones(5000, 5000)
startTime1 = time.time()
for i in range(100):
    C = torch.matmul(A, B)
endTime1 = time.time()
print('cpu计算总时长:', round((endTime1 - startTime1) * 1000, 2), 'ms')