现在pytorch主流的分布式训练库是DistributedDataParallel，它比Dataparallel库要快，而且前者能实现多机多卡后者只能单机多卡。本文是在单机多卡的环境下执行的分布式训练。
1. main.py（开启多进程）
首先用torch.multiprocess的spawn库来自动开启多进程进行分布式训练，每个子进程自动对应一个GPU和一个DDP训练的模块，这样就不需要自己手动执行多次main.py。
1）命令行参数：
这里省略了args的参数配置，可以根据自己情况设定，比如args.distributed_training指定所使用的GPU的个数。但命令行参数中必须有的是指定GPU索引， 例如：
CUDA_VISIBLE_DEVICES=0,1 python3 main.py --distributed_training 2这里CUDA_VISIBLE_DEVICES=0,1 指定了使用第0块和第1块GPU，如果不指定的话会默认使用第0块，当每个用服务器的人都不指定的话就会都用第0块，这就会造成out of memory，所以指定可见的GPU块号是非常重要的。
2）固定random seeds：
训练开始前要先固定numpy和torch的random seed，如果使用CUDA还要固定torch.manual_seed等，尤其设定cudnn.deterministic=True，这些是为了消除计算过程中的随机性，从而避免在参数相同、训练次数相同的情况下每次训练结果却都不同的问题。
# fix random seeds and fix CUDA seeds if using CUDA to avoid randomness of result
seed = args.random_seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if device != 'cpu':
# fix random seed and set deterministic=True
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
logger.info(f"finish setting seed:{seed}!")3）设置gpu_fraction：
如果有的GPU同时被其他人使用了，训练的时候很可能会报错out of memory，这是因为tensorflow会贪婪地尽可能占用更多的资源，甚至当资源不够用时会抢占未指定的GPU的资源。这时可以考虑设置gpu_fraction的方法，让tensorflow只占用每个指定GPU的部分资源，有时是能避免out of memory的。注意tensorflow的版本不同，写法会有点不同，如果tensorflow 2.x版本用第一种写法会报错“module 'tensorflow' has no attribute 'GPUOptions”，简单来说就是要把tf换成tf.compat.v1。
# (optional) set GPU fraction in case of 'out of memory'
# tensorflow 1.x
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.7
session = tf.Session(config=config)
# tensorflow 2.x
gpu_fraction = 0.7
gpu_options = tf.compat.v1.GPUOptions(allow_growth = True)
sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))4）main.py完整代码：
省略了args的配置。
import numpy as np
import random
import tensorflow as tf
import torch
import torch.multiprocessing as mp
from loguru import logger
if __name__ == "__main__":
logger.add(
"OurModel.log",
level='INFO',
rotation='12:00', # clean log files at 12:00 everyday
encoding='utf8',
backtrace=True,
diagnose=True,
enqueue=True
)
device = torch.device(
"cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
logger.info("device: {}".format(device))
# fix random seeds
seed = args.random_seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# fix CUDA seeds if using CUDA to avoid randomnes.
if device != 'cpu':
# fix random seed and set deterministic=True
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
logger.info(f"finish setting seed:{seed}!")
if args.do_train:
from do_train import Train
workers = Train.do_train
else:
from do_eval import Eval
workers = Eval.do_eval
if args.distributed_training:
# (optional) set GPU fraction in case of 'out of memory'
gpu_fraction = 0.7
# gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)
gpu_options = tf.compat.v1.GPUOptions(allow_growth = True)
sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))
# use parant queue to communicate between parant and children processes
queue = mp.Manager().Queue() # manager: parant process
kwargs = {'args': args}
spawn_context = mp.spawn(workers,
args=(args.distributed_training, # Total GPU number
kwargs, queue),
nprocs=args.distributed_training, # num of GPU per subprocess
join=False)
queue.put(True)
while not spawn_context.join():
pass
else:
assert False, 'Please set "--distributed_training 1" to use single gpu'2. do_train.py （子进程执行DDP）
整体流程：在do_train函数中初始化DDP配置，并将model放置到指定GPU上，dataset指定sampler=train_sampler，然后training epoch部分用train_sampler.set_epoch(i_epoch)给每个epoch sample dataset，最后destroy process group。
1）do_train 函数定义：
def do_train(rank, world_size, kwargs, queue)这里有四个参数，而main.py中mp.spawn的args参数列表中却只有三个参数，这是因为mp.spawn启动的worker函数自带第一个参数rank，对应的是CUDA_VISIBLE_DEVICES=0,1中等号右侧列表的下标，如rank=0对应"0,1"的第一个即0号GPU，如果CUDA_VISIBLE_DEVICES=1则rank=0对应1号GPU。worker函数第二个参数world_size也是mp.spawn一般都要使用的参数，表示使用的GPU个数；第三个参数kwargs传递args；queue参数由于在main.py中传入的变量是mp.Manager().Queue()，故特别用于父进程和子进程之间的通信。
2）DDP初始化：
这里的MASTER_ADDR和MASTER_PORT是用于设定端口给分布式进程之间通信的，如果端口被其他进程占用则会报错，这时需要换个端口号，端口号一般随机指定就好，例如29501。
dist.init_process_group用来初始化DDP中的每个进程，第一个参数"nccl"是通信协议，一般对于GPU分布式用“nccl”，CPU分布式用“gloo”。
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = args.master_port
# initialize the process group
dist.init_process_group("nccl", rank=rank, world_size=world_size)3）model 配置DDP
model = model(args).to(rank)
model = DDP(model, device_ids=[rank], find_unused_parameters=True)4）dataloader配置DDP
第一行是调用自己写的Dataset函数，获取预处理后的dataset；第二行是用DistributedSampler定义train_sampler；第三行调用的DataLoader与往常情况的区别除了sampler=train_sampler和batch_size=args.train_batch_size // world_size之外，还要注意是这里的shuffle要设成false（即保留默认值），num_workers也不需要设置。
train_dataset = Dataset(args, args.train_batch_size, to_screen=False)
train_sampler = DistributedSampler(train_dataset, shuffle=args.do_train)
train_dataloader = torch.utils.data.DataLoader(
# shuffle=False, num_workers=args.num_workers,
train_dataset, sampler=train_sampler,
batch_size=args.train_batch_size // world_size)
5）training epoch
train_sampler.set_epoch(i_epoch)用于sample每个epoch使用的dataset。
这里用tqdm加载dataset，可以实时看到训练dataset的进度（iter_bar）。
训练完后调用dist.barrier()阻塞分布式训练的每个进程，因为每个进程执行时间有先后，阻塞之后可以等所有进程都完成训练后再一起往下执行，从而实现进程同步。
6）do_train.py完整代码：
省略了train_one_epoch函数。
from loguru import logger
import os
import time
from functools import partial
import numpy as np
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
# from torch.utils.data import RandomSampler
from torch.utils.data.distributed import DistributedSampler
from data_preprocess.dataset import Dataset
from our_model import model
def learning_rate_decay(args, i_epoch, optimizer):
if i_epoch > 0 and i_epoch % 5 == 0:
for p in optimizer.param_groups:
p['lr'] *= 0.3
class Train:
def __init__(self):
pass
def do_train(rank, world_size, kwargs, queue):
# first args: rank is a default arg represents PID
# world_size: GPU_NUM
args = kwargs['args']
torch.cuda.set_device(rank)
# model
model = model(args).to(rank)
# pytorch_total_params = sum(p.numel() for p in model.parameters())
# logger.info(f'number of parameters: {pytorch_total_params}')
# distributed computing
if world_size > 0:
logger.info(f"Running DDP on rank {rank}.")
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = args.master_port
# initialize the process group
dist.init_process_group("nccl", rank=rank, world_size=world_size)
model = DDP(model, device_ids=[rank], find_unused_parameters=True)
# optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
if rank == 0 and world_size > 0:
receive = queue.get()
assert receive == True
if args.distributed_training:
dist.barrier() # block first-come process to synchronize all processes
# dataloader
logger.info(f"Loading dataset:{args.data_dir}")
train_dataset = Dataset(args, args.train_batch_size, to_screen=False)
train_sampler = DistributedSampler(train_dataset, shuffle=args.do_train)
train_dataloader = torch.utils.data.DataLoader(
train_dataset, sampler=train_sampler,
batch_size=args.train_batch_size // world_size)
# epochs
Max_Epoch = 1 if args.smoke_test else int(args.num_train_epochs)
if args.smoke_test:
logger.info("******Smoke Test******")
for i_epoch in range(Max_Epoch):
# learning rate
learning_rate_decay(args, i_epoch, optimizer)
logger.info(optimizer.state_dict()['param_groups'])
if rank == 0:
logger.info('Epoch: {}/{}'.format(i_epoch,
int(args.num_train_epochs)))
logger.info('Learning Rate = %5.8f' %
optimizer.state_dict()['param_groups'][0]['lr'])
train_sampler.set_epoch(i_epoch)
if rank == 0:
# iteration(progress) bar
iter_bar = tqdm(train_dataloader, desc='Iter (loss=X.XXX)')
else:
iter_bar = train_dataloader
train_one_epoch(model, iter_bar, optimizer,
rank, args, i_epoch, queue)
if args.distributed_training:
dist.barrier()
if args.distributed_training:
dist.destroy_process_group()
1532
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
