本文写于三年前,由于pytorch升级迭代,目前更推荐采用DDP实现方法,可看我新整理的一篇文章《Pytorch实现多GPU并行训练(DDP)_pytorch 多gpu并行训练-CSDN博客》
用pytorch进行多GPU训练,只需要学会把单卡训练的代码稍微改一下即可。不用弄得太麻烦。通过一个demo来做是最快入手的。
1. 要知道机器有几张卡:
nvidia-smi
2. 模型用DataParallel包装一下:
device_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 10卡机
model = torch.nn.DataParallel(model, device_ids=device_ids) # 指定要用到的设备
model = model.cuda(device=device_ids[0]) # 模型加载到设备0
3. 数据也指定设备:
X_train, y_train = X_train.cuda(device=device_ids[0]), y_train.cuda(device=device_ids[0])
这里只需要用device_ids[0]定义一个样式就好,不需要逐卡指定设备。但没这一步会报错。
4. 最后,来看一个完整demo,有注释的地方就是与单卡训练不一样的地方:
(代码来自https://www.cnblogs.com/fanghao/p/10334631.html,原代码在我的环境中运行有bug,下面的程序是修改版)
import torch
from torchvision import datasets, transforms
import torchvision
from tqdm import tqdm
device_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # 可用GPU
BATCH_SIZE = 64
transform = transforms.Compose([transforms.ToTensor()])
data_train = datasets.MNIST(root = "./data/",
transform=transform,
train=True,
download=True)
data_test = datasets.MNIST(root="./data/",
transform=transform,
train=False)
data_loader_train = torch.utils.data.DataLoader(dataset=data_train,
# 单卡batch size * 卡数
batch_size=BATCH_SIZE * len(device_ids),
shuffle=True,
num_workers=2)
data_loader_test = torch.utils.data.DataLoader(dataset=data_test,
batch_size=BATCH_SIZE * len(device_ids),
shuffle=True,
num_workers=2)
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.conv1 = torch.nn.Sequential(
torch.nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
torch.nn.ReLU(),
torch.nn.MaxPool2d(stride=2, kernel_size=2),
)
self.dense = torch.nn.Sequential(
torch.nn.Linear(14 * 14 * 128, 1024),
torch.nn.ReLU(),
torch.nn.Dropout(p=0.5),
torch.nn.Linear(1024, 10)
)
def forward(self, x):
x = self.conv1(x)
x = x.view(-1, 14 * 14 * 128)
x = self.dense(x)
return x
model = Model()
# 指定要用到的设备
model = torch.nn.DataParallel(model, device_ids=device_ids)
# 模型加载到设备0
model = model.cuda(device=device_ids[0])
cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())
from time import sleep
n_epochs = 50
for epoch in range(n_epochs):
running_loss = 0.0
running_correct = 0
print("Epoch {}/{}".format(epoch, n_epochs))
print("-"*10)
for data in tqdm(data_loader_train):
X_train, y_train = data
# 指定设备0
X_train, y_train = X_train.cuda(device=device_ids[0]), y_train.cuda(device=device_ids[0])
outputs = model(X_train)
_,pred = torch.max(outputs.data, 1)
optimizer.zero_grad()
loss = cost(outputs, y_train)
loss.backward()
optimizer.step()
running_loss += loss.data.item()
running_correct += torch.sum(pred == y_train.data)
testing_correct = 0
for data in data_loader_test:
X_test, y_test = data
# 指定设备1
X_test, y_test = X_test.cuda(device=device_ids[0]), y_test.cuda(device=device_ids[0])
outputs = model(X_test)
_, pred = torch.max(outputs.data, 1)
testing_correct += torch.sum(pred == y_test.data)
print("Loss is:{:.4f}, Train Accuracy is:{:.4f}%, Test Accuracy is:{:.4f}".format(torch.true_divide(running_loss, len(data_train)),
torch.true_divide(100*running_correct, len(data_train)),
torch.true_divide(100*testing_correct, len(data_test))))
torch.save(model.state_dict(), "model_parameter.pkl")
通过上面程序改造的方法,便可以把单卡训练改造成了多卡训练。当然,有些自定义的sampler,dataloader可能也要相应改之,具体问题具体处理。
我在另一个训练任务上,完全激活了10卡同时训练:10个2080ti,没有炫卡的意思(逃