Pytorch - GPU ID 指定 pytorch gpu 指定

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

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

巴特西

Pytorch - GPU ID 指定 pytorch gpu 指定

1. GPU ID 指定

2. torch.device

最新文章

热门文章