【AI实战】分布式训练：使用DistributedDataParallel实现单机多GPU并行训练resnet50模型

DistributedDataParallel

一种基于pytorch框架的分布式训练工具。

依赖包

import argparse
import time
import torch
import torchvision
from torch import distributed as dist
from torchvision.models import resnet18
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn

加载预训练模型

model_ft = models.resnet50(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层

使用DistributedDataParallel

net = model_ft
net.cuda()
net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = DDP(net, device_ids=[args.local_rank], output_device=args.local_rank)

自定义数据加载

class MyDataset(torch.utils.data.Dataset):
    
    def __init__(self, txt_path):
        
        im_list = []
        im_labels = []
        with open(txt_path, 'r') as files:
            for line in files:
                #/x/y/a.jpg 1
                #/x/y/b.jpg 2
                items = line.split()
                if len(items) != 2:
                    print(items)
                    continue
                im_list.append(items[0])
                im_labels.append(int(items[1]))
        self.imgs = im_list
        self.labels = im_labels
        
    def __len__(self):
        return len(self.imgs)
        
    def __getitem__(self, item):
        img_name = self.imgs[item]
        label = self.labels[item]
        
        def default_loader(path):
            with open(path, 'rb') as f:
                with Image.open(f) as img:
                    return img.convert('RGB')
        img = default_loader(img_name)

        try:
            img = data_tranforms(img)
        except:
            print("Cannot transform image: {}".format(img_name))
        return img, label

data_tranforms = transforms.Compose([
        transforms.Resize(224),
        
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
    ])

加载数据：

data_root = 'dataset'
trainset = MyDataset(txt_path='./data/Train.txt')
valset = MyDataset(txt_path='./data/Test.txt')

sampler = DistributedSampler(trainset)
train_loader = DataLoader(trainset,
                          batch_size=batch_size,
                          shuffle=False,
                          pin_memory=True,
                          sampler=sampler)
val_loader = DataLoader(valset,
                        batch_size=batch_size,
                        shuffle=False,
                        pin_memory=True)

分布式训练

criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=lr)

net.train()
for e in range(epochs):
    # DistributedSampler deterministically shuffle data
    # by seting random seed be current number epoch
    # so if do not call set_epoch when start of one epoch
    # the order of shuffled data will be always same
    sampler.set_epoch(e)
    for idx, (imgs, labels) in enumerate(train_loader):
        imgs = imgs.cuda()
        labels = labels.cuda()
        output = net(imgs)
        loss = criterion(output, labels)
        opt.zero_grad()
        loss.backward()
        opt.step()
        reduce_loss(loss, global_rank, world_size)
        if idx % 10 == 0 and global_rank == 0:
            print('Epoch: {} step: {} loss: {}'.format(e, idx, loss.item()))
net.eval()
with torch.no_grad():
    cnt = 0
    total = len(val_loader.dataset)
    for imgs, labels in val_loader:
        imgs, labels = imgs.cuda(), labels.cuda()
        output = net(imgs)
        predict = torch.argmax(output, dim=1)
        cnt += (predict == labels).sum().item()

if global_rank == 0:
    print('eval accuracy: {}'.format(cnt / total))

保存模型

## 保存模型 
path = './model/digit_classify-%s.pth' %(time.time())
torch.save(net.state_dict(), path)
print('*'*50)
print('data_tranforms', data_tranforms)
print('best model saved to ', path)
import shutil
shutil.copy(path,  './model/digit_classify.pth')

完整代码

import argparse
import time
import torch
import torchvision
from torch import distributed as dist
from torchvision.models import resnet18
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from PIL import Image
from torchvision import models, transforms
import torch.nn as nn


def reduce_loss(tensor, rank, world_size):
    with torch.no_grad():
        dist.reduce(tensor, dst=0)
        if rank == 0:
            tensor /= world_size

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', type=int, help="local gpu id")
args = parser.parse_args()

# python3 -m torch.distributed.launch --nproc_per_node=4 --master_port=29500 train.py
batch_size = 96
world_size = 4
epochs = 20
lr = 0.001
NUMCLASS = 11

dist.init_process_group(backend='nccl', init_method='env://')
torch.cuda.set_device(args.local_rank)
global_rank = dist.get_rank()
print('global_rank', global_rank)

#net = resnet18()
model_ft = models.resnet18(pretrained=True)
#model_ft = models.resnet50(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层
#model_ft = model_ft.to(device)
net = model_ft
net.cuda()
net = torch.nn.SyncBatchNorm.convert_sync_batchnorm(net)
net = DDP(net, device_ids=[args.local_rank], output_device=args.local_rank)


data_tranforms = transforms.Compose([
        transforms.Resize(448),
        
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
    ])

class MyDataset(torch.utils.data.Dataset):
    
    def __init__(self, txt_path):
        
        im_list = []
        im_labels = []
        with open(txt_path, 'r') as files:
            for line in files:
                #/x/y/a.jpg 1
                #/x/y/b.jpg 2
                items = line.split()
                if len(items) != 2:
                    print(items)
                    continue
                im_list.append(items[0])
                im_labels.append(int(items[1]))
        self.imgs = im_list
        self.labels = im_labels
        
    def __len__(self):
        return len(self.imgs)
        
    def __getitem__(self, item):
        img_name = self.imgs[item]
        label = self.labels[item]
        
        def default_loader(path):
            with open(path, 'rb') as f:
                with Image.open(f) as img:
                    return img.convert('RGB')
        img = default_loader(img_name)

        try:
            img = data_tranforms(img)
        except:
            print("Cannot transform image: {}".format(img_name))
        return img, label


data_root = 'dataset'
trainset = MyDataset(txt_path='./data/Train.txt')
valset = MyDataset(txt_path='./data/Test.txt')

sampler = DistributedSampler(trainset)
train_loader = DataLoader(trainset,
                          batch_size=batch_size,
                          shuffle=False,
                          pin_memory=True,
                          sampler=sampler)
val_loader = DataLoader(valset,
                        batch_size=batch_size,
                        shuffle=False,
                        pin_memory=True)



criterion = torch.nn.CrossEntropyLoss()
opt = torch.optim.Adam(net.parameters(), lr=lr)

net.train()
for e in range(epochs):
    # DistributedSampler deterministically shuffle data
    # by seting random seed be current number epoch
    # so if do not call set_epoch when start of one epoch
    # the order of shuffled data will be always same
    sampler.set_epoch(e)
    for idx, (imgs, labels) in enumerate(train_loader):
        imgs = imgs.cuda()
        labels = labels.cuda()
        output = net(imgs)
        loss = criterion(output, labels)
        opt.zero_grad()
        loss.backward()
        opt.step()
        reduce_loss(loss, global_rank, world_size)
        if idx % 10 == 0 and global_rank == 0:
            print('Epoch: {} step: {} loss: {}'.format(e, idx, loss.item()))
net.eval()
with torch.no_grad():
    cnt = 0
    total = len(val_loader.dataset)
    for imgs, labels in val_loader:
        imgs, labels = imgs.cuda(), labels.cuda()
        output = net(imgs)
        predict = torch.argmax(output, dim=1)
        cnt += (predict == labels).sum().item()

if global_rank == 0:
    print('eval accuracy: {}'.format(cnt / total))


## 保存模型 
path = './model/digit_classify-%s.pth' %(time.time())
torch.save(net.state_dict(), path)
print('*'*50)
print('data_tranforms', data_tranforms)
print('best model saved to ', path)
import shutil
shutil.copy(path,  './model/digit_classify.pth')