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')

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

参考:torch.cuda

Logo

为开发者提供学习成长、分享交流、生态实践、资源工具等服务,帮助开发者快速成长。

更多推荐