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()
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')
意思是先判断cuda是否存在，如果存在torch.device设为cuda，否则设置为cpu。
参考：torch.cuda
|