（转）PyTorch DDP模式单机多卡训练

一、启动训练的命令

python -m torch.distributed.launch --nproc_per_node=NUM_GPUS_YOU_HAVE train.py

其中torch.distributed.launch表示以分布式的方式启动训练，--nproc_per_node指定一共就多少个节点，可以设置成显卡的个数

二、启动之后每个进程可以自动获取到参数

import argparse
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", type=int,default=-1)
opt = parser.parse_args()
local_rank = opt.local_rank
print("local rank {}".format(local_rank))
assert torch.cuda.device_count() > opt.local_rank
torch.cuda.set_device(opt.local_rank)
device = torch.device('cuda', opt.local_rank)
dist.init_process_group(backend='nccl', init_method='env://')  # distributed backend
opt.world_size = dist.get_world_size()
print("world size {}".format(opt.world_size))
print("get rank {}".format(dist.get_rank()))

每个进程都能获取到local rank，local rank 表示的是进程的优先级，该优先级是自动分配的。world size 表示的一共运行的进程数和nproc_per_node设置的数值相对应。

 

三、正式开始DDP介绍训练模式设置

1.导入包

import torch
import torchvision

print("current torch version is {}".format(torch.__version__))
print("current torchvision version is {}".format(torchvision.__version__))

import sys
from models.resnet import *
from torchvision import datasets, transforms
import os
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
import time
import copy
from torch.nn import DataParallel
import argparse

2. 参数解读

parser = argparse.ArgumentParser()
parser.add_argument("--image_folder",type=str,default='/home/jl/datasets/oilrecognition',help='train and val folder path')
parser.add_argument("--local_rank", type=int,default=-1,help='DDP parameter, do not modify')#不需要赋值，启动命令 torch.distributed.launch会自动赋值
parser.add_argument("--distribute",action='store_true',help='whether using multi gpu train')
parser.add_argument("--distribute_mode",type=str,default='DDP',help="using which mode to ")
parser.add_argument('--epochs', type=int, default=20)
parser.add_argument('--batch_size', type=int, default=64, help='total batch size for all GPUs')
parser.add_argument("--save_path",type=str,default= "./save",help="the path used to save state_dict")
opt = parser.parse_args()

3. 初始化部分

if opt.distribute and opt.local_rank != -1:
     global device
     torch.cuda.set_device(opt.local_rank)
     torch.distributed.init_process_group(backend='nccl', init_method='env://')
     device = torch.device('cuda', opt.local_rank)

4. 数据部分

    data_dir = opt.image_folder
    image_datasets={}
    image_datasets['train'] = datasets.ImageFolder(os.path.join(data_dir, 'train'),data_transforms['train'])
    image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'),data_transforms['val'])
    
    word_size = dist.get_world_size()
    if opt.distribute and opt.local_rank != -1:
        train_sampler = torch.utils.data.distributed.DistributedSampler(image_datasets['train'],num_replicas = word_size,rank = opt.local_rank)
    else:
        train_sampler = None
    
    print("batch size is : {}".format(opt.batch_size))
    dataloaders = {}
    dataloaders['train'] = torch.utils.data.DataLoader(image_datasets['train'], batch_size=opt.batch_size,shuffle=(train_sampler is None), num_workers=4, pin_memory=True, sampler=train_sampler)
    dataloaders['val'] =  torch.utils.data.DataLoader(image_datasets['val'], batch_size=opt.batch_size,shuffle = False,num_workers=4)

5. 模型部分

 if opt.distribute and opt.local_rank != -1:
        model.to(device)
        model = DDP(model, device_ids=[opt.local_rank])

6.模型保存部分

if dist.get_rank()== -1 or 0:
        save_path = './oil_net.pt'
        torch.save(model.state_dict(), save_path)

作者：RunningJiang
链接：https://www.jianshu.com/p/7818b128b9cd
来源：简书
著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。