python指定Gpu

1、目前主流方法：.to(device)方法 （推荐）

import torch
import time
 

#1.通常用法
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
data = data.to(device)
model = model.to(device)
'''
    1.先创建device对象
    2.to函数指定数据或者模型放到哪里
'''

#2.将构建的tensor或者模型放到指定设备上（GPU）
torch.device('cuda',0) #这里的0指的是设备的序号
torch.device('cuda:0')


#3.例子 cpu转到GPU
s = time.time()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
a = torch.rand([1024, 1024, 10]).to(device)
print('time:{:6.3f}'.format(time.time()-s)) # 输出: time: 0.087

#4.例子2 Tensor GPU 转到cpu运行

 predict_boxes = predictions["boxes"].to("cpu").numpy()
 predict_classes = predictions["labels"].to("cpu").numpy()
 predict_scores = predictions["scores"].to("cpu").numpy()

.cuda()方法和.to(device)方法耗时基本差不多。但是.to(device)代码风格更容易后期修改。

2、to(device）需要注意的是：

使用GPU训练的时候，需要将Module对象和Tensor类型的数据送入到device。通常会使用 to.(device)。但是需要注意的是：

• 对于Tensor类型的数据，使用to.(device) 之后，需要接收返回值，返回值才是正确设置了device的Tensor。
• 对于Module对象，只用调用to.(device) 就可以将模型设置为指定的device。不必接收返回值。

# Module对象设置device的写法
model.to(device)

# Tensor类型的数据设置 device 的写法。
samples = samples.to(device)

当然model也可以接收返回值。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
img = img.to(device)

官网解释：

 

3、如果你有多个GPU

这样可以显示指定需要使用的计算资源，特别是有多个GPU的情况下。

那么可以参考以下代码：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = Model()

if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model，device_ids=[0,1,2])
    model.to(device)

1、torch.device用法总结_会飞的萨摩的博客-CSDN博客主要参考这两篇文章：

1、torch.device用法总结_会飞的萨摩的博客-CSDN博客

2、pytorch并行处理详解（多GPU，环境变量）_xys430381_1的博客-CSDN博客

--------------------------------旧材料-------------------------------------

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"] = "2"

见网址：（原）tensorflow中使用指定的GPU及GPU显存 - darkknightzh - 博客园

2. 使用函数 set_device

import torch
torch.cuda.set_device(id)

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

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

注：官方解释

如果服务器具有多个GPU，tensor.cuda()方法会将tensor保存到第一块GPU上，等价于tensor.cuda(0)。此时如果想使用第二块GPU，需手动指定tensor.cuda(1)，而这需要修改大量代码，很是繁琐。这里有两种替代方法：

• 一种是先调用t.cuda.set_device(1)指定使用第二块GPU，后续的.cuda()都无需更改，切换GPU只需修改这一行代码。
• 更推荐的方法是设置环境变量CUDA_VISIBLE_DEVICES，例如当export CUDA_VISIBLE_DEVICE=1（下标是从0开始，1代表第二块GPU），只使用第二块物理GPU，但在程序中这块GPU会被看成是第一块逻辑GPU，因此此时调用tensor.cuda()会将Tensor转移至第二块物理GPU。CUDA_VISIBLE_DEVICES还可以指定多个GPU，如export CUDA_VISIBLE_DEVICES=0,2,3，那么第一、三、四块物理GPU会被映射成第一、二、三块逻辑GPU，tensor.cuda(1)会将Tensor转移到第三块物理GPU上

设置CUDA_VISIBLE_DEVICES有两种方法，一种是在终端命令行中CUDA_VISIBLE_DEVICES=0,1 python main.py，一种是在程序中import os;os.environ["CUDA_VISIBLE_DEVICES"] = "2"。如果使用IPython或者Jupyter notebook，还可以使用%env CUDA_VISIBLE_DEVICES=1,2来设置环境变量

3.指定cuda的函数：使用torch.cuda.device，来指定默认使用哪一块GPU，或使用torch.set_default_tensor_type使程序默认使用GPU，不需要手动调用cuda。

# 如果未指定使用哪块GPU，默认使用GPU 0
x = t.cuda.FloatTensor(2, 3)
# x.get_device() == 0
y = t.FloatTensor(2, 3).cuda()
# y.get_device() == 0

# 指定默认使用GPU 1
with t.cuda.device(1):    
    # 在GPU 1上构建tensor
    a = t.cuda.FloatTensor(2, 3)

    # 将tensor转移至GPU 1
    b = t.FloatTensor(2, 3).cuda()
    print(a.get_device() == b.get_device() == 1 )

    c = a + b
    print(c.get_device() == 1)

    z = x + y
    print(z.get_device() == 0)

    # 手动指定使用GPU 0
    d = t.randn(2, 3).cuda(0)
    print(d.get_device() == 2)

t.set_default_tensor_type('torch.cuda.FloatTensor') # 指定默认tensor的类型为GPU上的FloatTensor
a = t.ones(2, 3)
a.is_cuda

4. tensor.cuda和variable.cuda,都会返回一个新对象，这个新对象的数据已转移至GPU，而之前的tensor/variable的数据还在原来的设备上（CPU）。但是nn.MOdule下的cuda()方法，会将所有数据都迁移至GPU，并返回自己。所以module = module.cuda()和module.cuda()效果相同。通过is_cuda()判断tensor对象是否在cuda上。

tensor = t.Tensor(3, 4)
# 返回一个新的tensor，保存在第1块GPU上，但原来的tensor并没有改变
tensor.cuda(0)
tensor.is_cuda # False

# 不指定所使用的GPU设备，将默认使用第1块GPU
tensor = tensor.cuda()
tensor.is_cuda # True

module = nn.Linear(3, 4)
module.cuda(device = 1)
module.weight.is_cuda # True

***后面是常用的方式****

5.判断支持cuda，自动决定是在GPU还是CPU上运行。

# 在不支持CUDA的机器下，下一步还是在CPU上运行
device = t.device("cuda:0" if t.cuda.is_available() else "cpu")
x = x.to(device)
y = y.to(x.device)
z = x+y

此外，还可以使用上面的tensor.cuda() 的方式将tensor拷贝到gpu上，但是这种方式不太推荐。

6.指定cuda的函数其他方式

>>> torch.zeros([2, 4], dtype=torch.int32)
tensor([[ 0,  0,  0,  0],
        [ 0,  0,  0,  0]], dtype=torch.int32)
>>> cuda0 = torch.device('cuda:0')
>>> torch.ones([2, 4], dtype=torch.float64, device=cuda0)
tensor([[ 1.0000,  1.0000,  1.0000,  1.0000],
        [ 1.0000,  1.0000,  1.0000,  1.0000]], dtype=torch.float64, device='cuda:0')

7.把数据从CPU迁移到GPU时，可以用.cuda()方法，也可以用.to(device)方法。示例如下。

.cuda()方法

import torch
import time
 
t = time.time()
b = torch.rand([1024, 1024, 10])
b = b.cuda()
print('time:{:6.3f}'.format(time.time() - t))  # 输出: time: 0.084

.to(device)方法 （推荐）

import torch
import time
 
s = time.time()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
a = torch.rand([1024, 1024, 10]).to(device)
print('time:{:6.3f}'.format(time.time()-s)) # 输出: time: 0.087

.cuda()方法和.to(device)方法耗时基本差不多。但是.to(device)代码风格更容易后期修改。

8.另外，值得注意到的是，张量在GPU和CPU之间的迁移不是in-place操作（返回新对象），而模型在GPU和CPU之间的迁移是in-place操作（返回原来的对象）。示例如下。

import torch
import torch.nn as nn
 
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
x_cpu = torch.randn(2, 3)
print('x_cpu id: ', id(x_cpu))
 
x_gpu = x_cpu.to(device)
print('x_gpu id: ', id(x_gpu))
 
 
net_cpu = nn.Linear(32, 10)
print('net_cpu id: ', id(net_cpu))
 
net_gpu = net_cpu.to(device)
print('net_gpu id: ', id(net_gpu))

运行结果：