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"
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
原文地址:https://www.cnblogs.com/jfdwd/p/11434332.html
时间: 2024-09-29 19:28:18