pytorch GPU torch.nn.DataParallel (DP) 多卡并行
本文的参考资料：
- Pytorch——多卡GUP训练原理(torch.nn.DataParallel) （博客园）
- [源码解析] PyTorch 分布式(2) ----- DataParallel(上)
- Pytorch的nn.DataParallel（知乎）
DataParallel的流程
DataParallel从流程上来看，是通过将一个mini-batch的数据加载到主线程上，然后数据均分到各个GPU中来工作。其流程包括以下几个步骤：
- 把模型从内存加载到GPU_0 (master GPU，主GPU) 中，把mini-batch数据从内存传输到主GPU。把模型分发(scatter)给各个GPU（也就是模型复制到各个GPU中）。将一个mini-batch的数据均分为多个sub-mini-batch，然后分发给各个GPU。
- 各个GPU独自完成sub-mini-batch的前向传播，计算得到output，并把output传递给主GPU。（PyTorch 使用多线程来并行前向传播，每个 GPU 在单独的线程上将针对各自的输入数据独立并行地进行 forward 计算。）
- 主GPU 收集 (gather) 各个GPU传递过来的output，并计算loss。即通过将网络输出与批次中每个元素的真实数据标签进行比较来计算损失函数值。
- 把计算得到的loss分发给其他GPU，在各个GPU之上运行后向传播，计算参数梯度。（之所以需要在各个GPU上计算梯度，是因为反向过程需要计算图的中间结果，GPU_0并没有其他GPU的中间计算结果。）
- 把各个GPU得到的梯度传递给主GPU，即在主GPU之上归并梯度：(1) 进行梯度下降，并更新主GPU上的模型参数；(2) 由于模型参数仅在主GPU上更新，而其他从属GPU此时并不是同步更新的，所以需要将更新后的模型参数复制到剩余的从属 GPU 中，以此来实现并行。
单GPU代码改成多GPU并行训练
第一步：设置GPU卡号
在程序的最开始的地方，添加下面这行代码。
os.environ["CUDA_VISIBLE_DEVICES"] = '2, 3' # 2,3是GPU的序号，根据需要进行修改。
上面这行代码设定了对这个程序而言可见的只有2和3号卡，和其他的卡没有关系，这是物理上的号卡，逻辑上来说其实是对应0和1号卡，即device_ids[0]对应的就是第2号卡，device_ids[1]对应的就是第3号卡。
上面这行代码一定要放在访问所有显卡的代码之前，否则无效。一般来讲，大家习惯放在import完package之后，也有人放在所有import之前。
第二步：把模型加载到多张GPU上
设置device：
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
这里的cuda:0对应的是上面第一步设置的第2号卡。
然后用 nn.DataParallel 封装模型；再用 mode.to(device) 将模型放于多块 GPUs 上：
if torch.cuda.device_count() > 1: model = nn.DataParallel(model) model.to(device)
第三步：保存模型
torch.save(net.module.state_dict(), path)
我习惯把模型先迁移到CPU，再进保存到硬盘：
torch.save(net.module.cpu().state_dict(), path)
加载模型的话，正常加载就可以了，例如：
net=Resnet18() net.load_state_dict(torch.load(path))
如果需要再加载到nn.DataParallel，再使用nn.DataParallel(net)进行包裹。
另外，有些博客文章中说需要把optimizer也放到nn.Dataparallel中，但是pytorch官方给的示例中并没有这样做。我做了几个小实验，发现也不需要把optimizer也放到nn.Dataparallel中。
实例
实例一：pytorch官方给的示例
import torch
import os
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
os.environ["CUDA_VISIBLE_DEVICES"] = '2, 3'
# Parameters and DataLoaders
input_size = 5
output_size = 2
batch_size = 30
data_size = 100
# Device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# Dummy Dataset
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
# dataloader
rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
batch_size=batch_size, shuffle=True)
# demo model - linear operation
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("\tIn Model: input size", input.size(),
"output size", output.size())
return output
# 单机多卡
# 1. 采用 nn.DataParallel 封装模型；
# 2. 采用 mode.to(device) 将模型放于多块 GPUs 上.
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)
model.to(device)
# 模型运行
# 查看 input tensors 和 output tensors 的sizes.
for data in rand_loader:
input = data.to(device)
output = model(input)
print("Outside: input size", input.size(),
"output_size", output.size())
实例二
import os
from datetime import datetime
import argparse
import torchvision
import torchvision.transforms as transforms
import torch
import torch.nn as nn
class ConvNet(nn.Module):
def __init__(self, num_classes=10):
super(ConvNet, self).__init__()
self.layer1 = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(16),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.layer2 = nn.Sequential(
nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(32),
nn.ReLU(),
nn.MaxPool2d(kernel_size=2, stride=2))
self.fc = nn.Linear(7*7*32, num_classes)
def forward(self, x):
out = self.layer1(x)
out = self.layer2(out)
out = out.reshape(out.size(0), -1)
out = self.fc(out)
return out
def train(args):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = ConvNet()
if torch.cuda.device_count() > 1:
print("==> Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model.to(device)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().to(device)
optimizer = torch.optim.SGD(model.parameters(), 1e-4)
# Data loading code
train_dataset = torchvision.datasets.MNIST(root='./data',
train=True,
transform=transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=8)
start = datetime.now()
total_step = len(train_loader)
for epoch in range(args.epochs):
for i, (images, labels) in enumerate(train_loader):
images = images.to(device)
labels = labels.to(device)
# Forward pass
outputs = model(images)
loss = criterion(outputs, labels)
# Backward and optimize
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i + 1) % 100 == 0:
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch + 1,
args.epochs,
i + 1,
total_step,
loss.item())
)
print("Training complete in: " + str(datetime.now() - start))
def main():
os.environ["CUDA_VISIBLE_DEVICES"] = '2, 3'
parser = argparse.ArgumentParser()
parser.add_argument('--epochs', default=2, type=int, metavar='N',
help='number of total epochs to run')
parser.add_argument('--batch_size', default=100, type=int, metavar='N',
help='batch size')
args = parser.parse_args()
train(args)
if __name__ == '__main__':
main()
浙公网安备 33010602011771号