PyTorch 關于多 GPUs 時的指定使用特定 GPU.

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

如：

import torcha = torch.Tensor(2, 3)
if torch.cuda.is_available(): # 判斷是否支持 CUDAa.is_cuda  # Falsea = a.cuda()  # 放到 GPU 上a.is_cuda  # Truea.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

2. torch.device

如：

import torchdevice = 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

總結

以上是生活随笔為你收集整理的Pytorch - GPU ID 指定 pytorch gpu 指定的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。