pytorch调用不了多个gpu_Pytorch多个GPU(数据并行)与单个GPU并行使用,PyTorch,GPUDataParallelism,的...

1.划重点

模型放到一个GPU上运行

model.gpu()

tensor = my_tensor.gpu()

模型放在多个GPU上运行

上文中的

model.gpu()

默认只使用一个GPU,如果你有多个GPU的话,

model = nn.DataParallel(model)

注意 DataParallel并行计算只存在在前向传播

2.有例子

下面通过一个线性回归的例子来说明;一个输出通过线性变换得到一个结果

#包的导入

import torch

import torch.nn as nn

from torch.autograd import Variable

from torch.utils.data import Dataset, DataLoader

# Parameters and DataLoaders

input_size = 5

output_size = 2

batch_size = 30

data_size = 100

#创建类,获取随机数

class RandomDataset(Dataset):

def __init__(self, size, length):

self.len = length

self.data = torch.randn(length, size)

def __getitem__(self, index):

return self.data[index]

def __len__(self):

return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, 100),

batch_size=batch_size, shuffle=True)

#构建线性网络,仅用前向传播,注

class Model(nn.Module):

def __init__(self, input_size, output_size):

super(Model, self).__init__()

self.fc = nn.Linear(input_size, output_size)

def forward(self, input):

output = self.fc(input)

print(" In Model: input size", input.size(),

"output size", output.size())

return output

模型实例化和数据并行

首先,我们要多模型进行实例化然后检查是不是有多个GPUs,如果是的话就要先用nn.DataParallel语句,然后就可以调用model.gpu()将模型放到GPUs上面。如果只有一个GPU那就直接调用model.gpu()就可以了。

model = Model(input_size, output_size)

if torch.cuda.device_count() > 1:

print("Let's use", torch.cuda.device_count(), "GPUs!")

# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs

model = nn.DataParallel(model)

if torch.cuda.is_available():

model.cuda()

运行程序(GPU 仅有一个,2080Ti)

for data in rand_loader:

if torch.cuda.is_available():

input_var = Variable(data.cuda())

else:

input_var = Variable(data)

output = model(input_var)

print("Outside: input size", input_var.size(),

"output_size", output.size())

运行结果:

In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([30, 5]) output size torch.Size([30, 2])

Outside: input size torch.Size([30, 5]) output_size torch.Size([30, 2])

In Model: input size torch.Size([10, 5]) output size torch.Size([10, 2])

Outside: input size torch.Size([10, 5]) output_size torch.Size([10, 2])

你可能感兴趣的:(pytorch调用不了多个gpu_Pytorch多个GPU(数据并行)与单个GPU并行使用,PyTorch,GPUDataParallelism,的...)