【设置gpu设备】os.environ[‘CUDA_VISIBLE_DEVICES‘] 和 torch.cuda.set_device()

【设置gpu设备】os.environ[‘CUDA_VISIBLE_DEVICES‘] 和 torch.cuda.set_device()

1. 介绍

官方文档:当使用 PyTorch 进行深度学习训练时,通常需要使用 CUDA 加速计算。在使用 PyTorch 进行训练之前,需要确保已经正确设置了可见的 GPU 设备,并且已经初始化了 CUDA 环境。

2. 方法

2.1 方法1:os.environ[‘CUDA_VISIBLE_DEVICES‘] (推荐)

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

os.environ[‘CUDA_VISIBLE_DEVICES’] 是一个环境变量,可以通过设置它来限制程序所能看到的可用 GPU 设备列表,从而确保程序只使用指定的 GPU 设备。

  • 设置该环境变量可以使用 os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0,1’,其中的 ‘0,1’ 表示程序只能看到编号为 0 和 1 的 GPU 设备。

2.2 方法2:torch.cuda.set_device(0)

import torch
torch.cuda.set_device(0)

torch.cuda.set_device() 则是一个 PyTorch 提供的函数,用于将程序的运行环境切换到指定的 GPU 设备上。可以使用 torch.cuda.set_device(0) 将程序的运行环境切换到编号为 0 的 GPU 设备上。

2.3 说明

当你发现你用方式一没有办法使用你指定的GPU,可以用方式二。

  • data.cuda(gpu_id, non_blocking=non_blocking)

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

3. 参考

【1】https://blog.csdn.net/weixin_46141646/article/details/129052145

  • 6
    点赞
  • 13
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

秋冬无暖阳°

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值