Pytorch - GPU ID 指定 pytorch gpu 指定

谁践踏了优雅 2023-06-04 15:59 46阅读 0赞

PyTorch 关于多 GPUs 时的指定使用特定 GPU.

PyTorch 中的 Tensor，Variable 和 nn.Module(如 loss，layer和容器 Sequential) 等可以分别使用 CPU 和 GPU 版本，均是采用 `.cuda()` 方法.

如：

import torch a = torch.Tensor(2, 3) if torch.cuda.is_available(): # 判断是否支持 CUDA a.is_cuda # False a = a.cuda() # 放到 GPU 上 a.is_cuda # True a.get_device() # 默认使用 GPU-0

采用 `.cuda()` 方法默认使用的是 GPU-0，等价于 `.cuda(0)`.

## 1. GPU ID 指定 ##

当需要指定使用多张 GPUs 中的特定 GPU 时，可以采用 `.cuda(1)` 方法，但需要对大量的 Tensor、Variable等进行修改.

参考网络上的方法，替代方案主要有：

\[1\] - 使用 `torch.cuda.set_device(id)` 函数

import torch torch.cuda.set_device(id) # id=0, 1, 2 等

\[2\] - 采用类似 Tensorflow 指定 GPU 的方式，使用 `CUDA_VISIBLE_DEVICES`

# 使用终端命令行运行的 GPU 指定方式 CUDA_VISIBLE_DEVICES=1 python python_script.py

# 在 python 脚本中的 GPU 指定方式 import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"

> 参考：[PyTorch中使用指定的GPU][PyTorch_GPU]

## 2. torch.device ##

如：

import torch device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") a_cpu = torch.Tensor(2, 3) # tensor([[0.0000e+00, 0.0000e+00, 1.2771e-40], # [9.0079e+15, 1.6751e-37, 2.9775e-41]]) a_cpu.device # device(type='cpu') a_cuda = a.to(device) # tensor([[-2.1800e-01, 4.5737e-41, 2.3351e-37], # [ 0.0000e+00, 4.4842e-44, 0.0000e+00]], device='cuda:0') a_cuda.device # device(type='cuda', index=0) # 查看当前 gpu id curr_gpuid = torch.cuda.current_device() # 0

![1432127-20190830121114142-1827886879.jpg][]

转载于:https://www.cnblogs.com/jfdwd/p/11434332.html

[PyTorch_GPU]: https://www.cnblogs.com/darkknightzh/p/6836568.html
[1432127-20190830121114142-1827886879.jpg]: /images/20230601/96a61423630a4ecda294fb20907b0502.png