一、單機單GPU
1、set current device (gpu id)
# the first methord
CUDA_VISIBLE_DEVICES=gpi_id python XXX.py
# the second methord
torch.cuda.set_device(gpu_id)
2、put the object(tensor,variable,model…) to GPU memory
# 1、Tensor
ten1 = torch.FloatTensor(2).cuda()
# 2、Variable
ten1 = torch.FloatTensor(2)
#first variable,then cuda
V1_cpu = autograd.Variable(ten1)
V1 = V1_cpu.cuda()
# first cuda,then variable
ten1_cuda = ten1.cuda()
V2 = autograd.Variable(ten1_cuda)
3、Returns a copy of this object in CPU memory,then you can interact with Numpy
V1_cpu = V1.cpu()
V2_cpu = V2.cpu()
二、單機多GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
- The torch.device contains a device type (‘cpu’ or ‘cuda’) and optional device ordinal for the device type.
model = Model(input_size,output_size) # init the model
if torch.cuda.device_count() > 1: # Returns the number of GPUs available
model = nn.DataParallel(model)
model.to(device)
NB:The batch size should be larger than the number of GPUs used.
the link will help you
DATA PARALLELISM
pytorch 多GPU訓練總結(DataParallel的使用)
and you can realize how to select a device in th cross-gpu operations.
CUDA SEMANTICS