python specified GPU

1. The current mainstream method: .to(device) method (recommended)

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() method and .to(device) method take almost the same time. However, the .to(device) code style is easier to modify later.

2. To (device) needs to pay attention to:

When using GPU training, Module objects and Tensor type data need to be sent to the device. Usually to.(device) is used. But what needs to be noted is:

• For Tensor type data, after using to.(device), you need to receive the return value. The return value is the Tensor with the device correctly set.
• For Module objects, just call to.(device) to set the model to the specified device. It is not necessary to receive a return value.

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

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

Of course, model can also receive return values.

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

Official website explanation:

 

3. If you have multiple GPUs

This allows you to display and specify the computing resources that need to be used, especially when there are multiple GPUs.

Then you can refer to the following code:

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. Summary of usage of torch.device_Flying Satsuma's Blog-CSDN BlogMainly refer to these two articles:

1,torch.device Usage Summary_Flying Satsuma's Blog-CSDN Blog

2,Detailed explanation of pytorch parallel processing (multiple GPUs, environment variables)_xys430381_1's blog-CSDN blog

--------------------------------Old material---------------- ---------------------

PyTorch uses GPU starting from 0 by default. If GPU0 is running the program, you need to specify other GPUs.

There are two ways to specify the GPU to be used.

1. Similar to how tensorflow specifies GPU, use CUDA_VISIBLE_DEVICES.

1.1 Set directly in the terminal:

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 Settings in python code:

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"

See URL: (Original) Use the specified GPU and GPU memory in tensorflow - darkknightzh - Blog Park

2. Use the function set_device

import torch
torch.cuda.set_device(id)

This function can be found in pytorch-master\torch\cuda\__init__.py.

However, it is officially recommended to use CUDA_VISIBLE_DEVICES, and it is not recommended to use the set_device function.

Note: official explanation

If the server has multiple GPUs, the tensor.cuda() method will save the tensor to the first GPU, which is equivalent to tensor.cuda(0). If you want to use the second GPU at this time, you need to manually specify tensor.cuda(1), and this requires modifying a lot of code, which is very cumbersome. Here are two alternatives:

• One is to callt.cuda.set_device(1) first to specify the use of the second GPU. Subsequent.cuda() does not need to be changed. To switch GPUs, you only need to modify this line of code.
• The more recommended method is to set the environment variableCUDA_VISIBLE_DEVICES, for example, whenexport CUDA_VISIBLE_DEVICE=1 (the subscript starts from 0, 1 represents the second GPU), Only the second physical GPU is used, but this GPU will be regarded as the first logical GPU in the program, so calling tensor.cuda() at this time will transfer the Tensor to the second physical GPU. CUDA_VISIBLE_DEVICESYou can also specify multiple GPUs, such as export CUDA_VISIBLE_DEVICES=0,2,3, then the first, third, and fourth physical GPUs will be mapped to the first, second, and third logical GPUs ,tensor.cuda(1)will transfer the Tensor to the third physical GPU

There are two methods for settingCUDA_VISIBLE_DEVICES,one is in the terminal command lineCUDA_VISIBLE_DEVICES=0,1 python main.py, and the other is In programimport os;os.environ["CUDA_VISIBLE_DEVICES"] = "2". If you use IPython or Jupyter notebook, you can also use %env CUDA_VISIBLE_DEVICES=1,2 to set environment variables

3. Specify the cuda function: Use torch.cuda.device to specify which GPU to use by default, or use torch.set_default_tensor_type Make the program use GPU by default, no need to call cuda manually.

# 如果未指定使用哪块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') # Specify the default tensor type as FloatTensor on the GPU
a = t.ones(2, 3)
a.is_cuda

4. Both tensor.cuda and variable.cuda will return a new object. The data of this new object has been transferred to the GPU, while the previous tensor/variable data is still on the original device (CPU). However, the cuda() method under nn.MOdule will migrate all data to the GPU and return itself. So module = module.cuda() and module.cuda() have the same effect. Use is_cuda() to determine whether the tensor object is on cuda.

tensor = t.Tensor(3, 4)
# Return a new tensor, which is saved on the first GPU, but the original tensor has not changed
tensor.cuda(0)
tensor.is_cuda # False

# If the GPU device used is not specified, the first GPU will be used by default
tensor = tensor.cuda()
tensor.is_cuda # True

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

***The following are commonly used methods****

5. Determine whether CUDA is supported and automatically decide whether to run on GPU or CPU.

# In a machine that does not support CUDA, the next step is to run it on the 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

In addition, can also use the above tensor.cuda() method to copy the tensor to the GPU, but this method is not recommended.

6. Other ways to specify cuda functions

>>> 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. When migrating data from CPU to GPU, you can use the .cuda() method or the .to(device) method. An example is as follows.

.cuda() method

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) method (recommended)

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() method and .to(device) method take almost the same time. However, the .to(device) code style is easier to modify later.

8. In addition, it is worth noting that the migration of tensors between GPU and CPU is not an in-place operation (returning a new object), and The migration of the model between GPU and CPU is an in-place operation (returning the original object). Examples are as follows.

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))

operation result: