PyTorch如何获得显卡 Compute Compatibility

前言

什么是Compute compatibility（下文简称CC）？它是NVIDIA为各代显卡（包括Jetson）设定的一个值，和算力没关系，直译过来叫“计算兼容性”，通常更高的CC可以跑更多的网络结构，硬编码质量也更高。

正文

老规矩，直接上代码：

import torch
cc = torch.cuda.get_device_capability()
cc = f'{cc[0]}{cc[1]}'
print(cc)

虽然很简单，还是解读一下吧

torch.cuda.get_device_capability()

输出的是一个tuple，含两个元素，都是int。tuple里有int不能''.join()，所以下一行把它们放进一个字符串里（有更好的方法欢迎评论）。你当然也可以直接用tuple获取前一位或后一位进行大小比较。

比如我是P100显卡，torch.cuda.get_device_capability()的输出就是(6, 0)，整个脚本的结果是'60'。

补充

想要知道自己显卡的CC可以直接去英伟达官网查：CUDA GPUs | NVIDIA Developer。

说点小规律：Amperer, Volta, Pascal对应的第一位CC分别是8、7、6；Tesla x100 GPU第二位CC都是0。