27 lines
639 B
Python
27 lines
639 B
Python
import time
|
|
import torch
|
|
import sys
|
|
|
|
import os
|
|
|
|
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
|
|
|
|
# 测试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')
|
|
|
|
sys.exit(1)
|
|
|
|
# # 测试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') |