#方式一 #import os # os.environ['CUDA_VISIBLE_DEVICES'] = '5' #方式二 #import torch #torch.cuda.set_device(5)
官方文档:当使用 PyTorch 进行深度学习训练时,通常需要使用 CUDA 加速计算。在使用 PyTorch 进行训练之前,需要确保已经正确设置了可见的 GPU 设备,并且已经初始化了 CUDA 环境。
os.environ['CUDA_VISIBLE_DEVICES']
是一个环境变量,可以通过设置它来限制程序所能看到的可用 GPU 设备列表,从而确保程序只使用指定的 GPU 设备。设置该环境变量可以使用 os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
,其中的 '0,1'
表示程序只能看到编号为 0 和 1 的 GPU 设备。
torch.cuda.set_device()
则是一个 PyTorch 提供的函数,用于将程序的运行环境切换到指定的 GPU 设备上。可以使用 torch.cuda.set_device(0)
将程序的运行环境切换到编号为 0 的 GPU 设备上。
当你发现你用方式一没有办法使用你指定的GPU,可以用方式二。
在使用:文章来源:https://www.toymoban.com/news/detail-403705.html
data.cuda(gpu_id, non_blocking=non_blocking)
如果报错,可以改为data.cuda()文章来源地址https://www.toymoban.com/news/detail-403705.html
到了这里,关于os.environ[‘CUDA_VISIBLE_DEVICES‘] 和torch.cuda.set_device()的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!