torch.cuda common instructions

CUDA (Compute Unified Device Architecture) is a computing platform launched by graphics card manufacturer NVIDIA . CUDA™ is a general-purpose parallel computing architecture introduced by NVIDIA that enables GPUs to solve complex computing problems. It includes the CUDA instruction set architecture ( ISA ) and the parallel computing engine inside the GPU. Developers can use the C language to write programs for the CUDA™ architecture, and the written programs can run at ultra-high performance on CUDA™-enabled processors.

Pytorch allows models and data to be loaded onto the GPU through cuda instructions. Common instructions are as follows:

1. torch.cuda.is_available()

Is cuda available  

import torch
print(torch.cuda.is_available())

2. torch.cuda.device_count()

View the number of GPUs

import torch
print(torch.cuda.device_count())

3. torch.cuda.get_device_name()

View DEVICE (GPU) name

import torch
print(torch.cuda.get_device_name())

4. torch.cuda.current_device()

Check the serial number of the GPU currently in use

import torch
print(torch.cuda.current_device())

5. torch.cuda.set_device()

designated card

torch.cuda.set_device(gpu_id) #单卡
torch.cuda.set_device('cuda:'+str(gpu_ids)) #多卡

Specify only the main graphics card, as follows:

import torch

torch.cuda.set_device(1)
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

Specify a specific graphics card, as follows:

import torch
import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'

torch.cuda.set_device(1)
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

Using torch.cuda.set_device() can load the model and data to the corresponding GPU more conveniently, just add a line of code before defining the model, but this way of writing has a low priority, if model.cuda() specifies parameters, then torch.cuda.set_device() will fail, and the official documentation of pytorch clearly states that users are not recommended to use this method.

6.  .cuda()

The specified model and data are loaded to the corresponding GPU. Taking net.cuda() as an example, the loading method is:

net.cuda(gpu_id) # 输入参数为int类型，只能指定一张显卡

net.cuda('cuda:'+str(gpu_ids)) #输入参数为str类型，可指定多张显卡

Specify only the main graphics card, as follows:

import torch
import os 
 
x = torch.tensor([[1,2,3],[4,5,6]]).cuda()
print(x.device)

Specify a specific graphics card, as follows:

import torch
import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2'

x = torch.tensor([[1,2,3],[4,5,6]]).cuda('cuda:1')
print(x.device)