os.environ[‘CUDA_VISIBLE_DEVICES‘] 和torch.cuda.set_device()

最新推荐文章于 2025-05-14 22:27:35 发布

weixin_46141646

最新推荐文章于 2025-05-14 22:27:35 发布

阅读量3.8k

点赞数 2

文章标签：深度学习人工智能 python

本文链接：https://blog.youkuaiyun.com/weixin_46141646/article/details/129052145

版权

在PyTorch的深度学习训练中，设置CUDA设备是关键步骤。文章介绍了两种方法：一是通过环境变量`CUDA_VISIBLE_DEVICES`指定GPU，如`os.environ[CUDA_VISIBLE_DEVICES]=5`；二是使用PyTorch的`torch.cuda.set_device(5)`函数切换到特定GPU。若遇到问题，可以尝试直接使用`data.cuda()`迁移数据到GPU。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

#方式一 
#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，可以用方式二。

在使用:

data.cuda(gpu_id, non_blocking=non_blocking)

如果报错，可以改为data.cuda()