使用pytorch代码测试:
# 矩阵乘法的TFLOPS
import torch
from torch.utils import benchmark
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
typ = torch.float16
n = 1024 * 16
a = torch.randn(n, n).type(typ).to(device)
b = torch.randn(n, n).type(typ).cuda()
t = benchmark.Timer(
stmt='a @ b',
globals={'a':a, 'b':b}
)
x = t.timeit(50)
print(2*n**3 / x.median / 1e12)
结果:
20.800833609440435
我的电脑是笔记本,配置如下:
CPU:R7 4800H
GPU:NVIDIA RTX 2060
使用cuda 11.7,3090 Ti能跑到85(理论160),V100 16GB能跑86(理论125),A100 80GB能跑到231(理论312)。
有用请点个赞!!
本站所有文章均为原创,欢迎转载,请注明文章出处:https://blog.csdn.net/weixin_45092662。百度和各类采集站皆不可信,搜索请谨慎鉴别。技术类文章一般都有时效性,本人习惯不定期对自己的博文进行修正和更新,因此请访问出处以查看本文的最新版本。