PyTorch 中如何指定GPU

本文介绍了在PyTorch中指定GPU的方法,包括通过设置环境变量CUDA_VISIBLE_DEVICES和使用torch.cuda.set_device函数,并推荐了首选方案。此外,还说明了如何将张量和函数转移到GPU上。

PyTorch默认使用从0开始的GPU,如果GPU0正在运行程序,需要指定其他GPU。

有如下两种方法来指定需要使用的GPU。

1. 类似tensorflow指定GPU的方式,使用CUDA_VISIBLE_DEVICES。

1.1 直接终端中设定:

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代码中设定:

## 建议使用该方法
import os 
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

见网址:http://www.cnblogs.com/darkknightzh/p/6591923.html

2. 使用函数 set_device

##不建议使用该方法
import torch 
torch.cuda.set_device(1)

该函数见 pytorch-master\torch\cuda\__init__.py。

不过官方建议使用CUDA_VISIBLE_DEVICES,不建议使用 set_device 函数。

以上引用自:https://www.cnblogs.com/darkknightzh/p/6836568.html

2.如何将tensor或函数放在GPU上

 直接使用tensor.cuda() 或 func().cuda()即可。不能直接对类或类对象执行".cuda()"

在VSCode中安装支持GPUPyTorch,可按以下步骤进行: ### 1. 确认CUDA版本 首先需要确认系统中安装的CUDA版本,因为PyTorchCUDA版本有一定要求。可以通过以下命令查看CUDA版本: ```bash nvcc --version ``` 如果系统未安装CUDA,需要从NVIDIA官方网站下载并安装适合自己显卡的CUDA版本。 ### 2. 创建虚拟环境 为了避免不同项目之间的依赖冲突,建议创建一个虚拟环境。可以使用`conda`或`venv`来创建虚拟环境。 #### 使用conda创建虚拟环境 ```bash conda create -n pytorch_gpu python=3.8 conda activate pytorch_gpu ``` #### 使用venv创建虚拟环境 ```bash python -m venv pytorch_gpu source pytorch_gpu/bin/activate # Linux/Mac pytorch_gpu\Scripts\activate # Windows ``` ### 3. 安装支持GPUPyTorch 根据之前确认的CUDA版本,从PyTorch官方网站选择合适的安装命令。例如,如果CUDA版本是11.3,可以使用以下命令安装: ```bash pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113 ``` ### 4. 在VSCode中配置Python解释器 打开VSCode,打开项目文件夹。按下`Ctrl + Shift + P`(Windows/Linux)或`Cmd + Shift + P`(Mac),输入`Python: Select Interpreter`,选择之前创建的虚拟环境中的Python解释器。 ### 5. 验证安装 在VSCode中创建一个Python文件,输入以下代码来验证PyTorch是否支持GPU: ```python import torch print(torch.cuda.is_available()) ``` 运行代码,如果输出`True`,则说明PyTorch已经成功安装并且支持GPU
评论
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

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

余额充值