torch.cuda

• 该包增加了对CUDA张量类型的支持，实现了与CPU张量相同的功能，但使用GPU进行计算。
• 它是延迟的初始化，所以你可以随时导入它，并使用is_available()来确定系统是否支持CUDA。
• CUDA语义有关于使用CUDA的更多细节。

torch.cuda.is_available()

作用：返回bool值，指示当前CUDA是否可用。

torch.cuda.set_device(device)

• 作用：设置当前设备。不鼓励使用此功能函数。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。
• 参数：device(int)表示选择的设备。如果此参数为负，则此函数是无操作的。

class torch.cuda.device(idx)

• 作用：更改所选设备的上下文管理器。
• 参数：idx(int) 表示设备索引选择。如果这个参数是负的，则是无效操作。

torch.cuda.device_count()

• 作用：返回gpu数量。

torch.cuda.get_device_name(0)

作用：返回gpu名字，设备索引默认从0开始。

torch.cuda.current_device()

作用：返回当前设备索引。

torch.device(‘cuda’)

作用：将数据转移到GPU。

torch.device(‘cpu’)

作用：将数据转移的CPU。

torch.device()表示torch.Tensor被分配到的设备对象，共有cpu和cuda两种，这里的cuda指的就是gpu，至于为什么不直接用gpu与cpu对应，是因为gpu的编程接口采用的是cuda。

例：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
1

意思是先判断cuda是否存在，如果存在torch.device设为cuda，否则设置为cpu。

参考

1. torch.cuda ↩