为什么torch.cuda.device的设备参数可能大于torch.cuda.device_count()？

Question

这是代码：

torch.cuda.device_count()
# output: 1
d2 = torch.cuda.device(2)
# d2.idx = 2

是否适用于在其他设备上运行的机型？

Answer 1

torch.cuda.device_count()返回可用的GPU数量，而torch.cuda.device(device)是context-manager。它有助于正确处理资源，即资源使用后的自动设置和释放。因此，您在此函数中传递的参数并不意味着您要切换/设置到该设备。为了设置当前设备，您应该使用torch.cuda.set_device(device)。虽然，根据官方文档here，

为了支持设备，不鼓励使用此函数的

。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。

要了解更多关于torch.cuda.device()及其工作原理的信息，请访问this官方pytorch讨论。