Pytorch GPU多卡并行訓(xùn)練實(shí)戰(zhàn)總結(jié)（附代碼）
大家好，我是DASOU；
今天分享給大家一份Pytorch?GPU多卡并行訓(xùn)練實(shí)戰(zhàn)細(xì)節(jié)總結(jié)
為什么要使用多GPU并行訓(xùn)練
簡單來說，有兩種原因：第一種是模型在一塊GPU上放不下，兩塊或多塊GPU上就能運(yùn)行完整的模型（如早期的AlexNet）。第二種是多塊GPU并行計(jì)算可以達(dá)到加速訓(xùn)練的效果。想要成為“煉丹大師“，多GPU并行訓(xùn)練是不可或缺的技能。
常見的多GPU訓(xùn)練方法：
1.模型并行方式：如果模型特別大，GPU顯存不夠，無法將一個(gè)顯存放在GPU上，需要把網(wǎng)絡(luò)的不同模塊放在不同GPU上，這樣可以訓(xùn)練比較大的網(wǎng)絡(luò)。（下圖左半部分）
2.數(shù)據(jù)并行方式：將整個(gè)模型放在一塊GPU里，再復(fù)制到每一塊GPU上，同時(shí)進(jìn)行正向傳播和反向誤差傳播。相當(dāng)于加大了batch_size。（下圖右半部分）
在pytorch1.7 + cuda10 + TeslaV100的環(huán)境下，使用ResNet34，batch_size=16, SGD對花草數(shù)據(jù)集訓(xùn)練的情況如下：使用一塊GPU需要9s一個(gè)epoch，使用兩塊GPU是5.5s， 8塊是2s。這里有一個(gè)問題，為什么運(yùn)行時(shí)間不是9/8≈1.1s ? 因?yàn)槭褂肎PU數(shù)量越多，設(shè)備之間的通訊會越來越復(fù)雜，所以隨著GPU數(shù)量的增加，訓(xùn)練速度的提升也是遞減的。
誤差梯度如何在不同設(shè)備之間通信？
在每個(gè)GPU訓(xùn)練step結(jié)束后，將每塊GPU的損失梯度求平均，而不是每塊GPU各計(jì)算各的。
BN如何在不同設(shè)備之間同步？
假設(shè)batch_size=2，每個(gè)GPU計(jì)算的均值和方差都針對這兩個(gè)樣本而言的。而BN的特性是：batch_size越大，均值和方差越接近與整個(gè)數(shù)據(jù)集的均值和方差，效果越好。使用多塊GPU時(shí)，會計(jì)算每個(gè)BN層在所有設(shè)備上輸入的均值和方差。如果GPU1和GPU2都分別得到兩個(gè)特征層，那么兩塊GPU一共計(jì)算4個(gè)特征層的均值和方差，可以認(rèn)為batch_size=4。注意：如果不用同步BN，而是每個(gè)設(shè)備計(jì)算自己的批次數(shù)據(jù)的均值方差，效果與單GPU一致，僅僅能提升訓(xùn)練速度；如果使用同步BN，效果會有一定提升，但是會損失一部分并行速度。
下圖為單GPU、以及是否使用同步BN訓(xùn)練的三種情況，可以看到使用同步BN（橙線）比不使用同步BN（藍(lán)線）總體效果要好一些，不過訓(xùn)練時(shí)間也會更長。使用單GPU（黑線）和不使用同步BN的效果是差不多的。
?兩種GPU訓(xùn)練方法：DataParallel 和 DistributedDataParallel：
DataParallel是單進(jìn)程多線程的，僅僅能工作在單機(jī)中。而DistributedDataParallel是多進(jìn)程的，可以工作在單機(jī)或多機(jī)器中。 DataParallel通常會慢于DistributedDataParallel。所以目前主流的方法是DistributedDataParallel。
pytorch中常見的GPU啟動方式：
注：distributed.launch方法如果開始訓(xùn)練后，手動終止程序，最好先看下顯存占用情況，有小概率進(jìn)程沒kill的情況，會占用一部分GPU顯存資源。
下面以分類問題為基準(zhǔn)，詳細(xì)介紹使用DistributedDataParallel時(shí)的過程:
首先要初始化各進(jìn)程環(huán)境：
def init_distributed_mode(args):
# 如果是多機(jī)多卡的機(jī)器，WORLD_SIZE代表使用的機(jī)器數(shù)，RANK對應(yīng)第幾臺機(jī)器
# 如果是單機(jī)多卡的機(jī)器，WORLD_SIZE代表有幾塊GPU，RANK和LOCAL_RANK代表第幾塊GPU
if'RANK'in os.environ and'WORLD_SIZE'in os.environ:
args.rank = int(os.environ["RANK"])
args.world_size = int(os.environ['WORLD_SIZE'])
# LOCAL_RANK代表某個(gè)機(jī)器上第幾塊GPU
args.gpu = int(os.environ['LOCAL_RANK'])
elif'SLURM_PROCID'in os.environ:
args.rank = int(os.environ['SLURM_PROCID'])
args.gpu = args.rank % torch.cuda.device_count()
else:
print('Not using distributed mode')
args.distributed = False
return
args.distributed = True
torch.cuda.set_device(args.gpu) # 對當(dāng)前進(jìn)程指定使用的GPU
args.dist_backend = 'nccl'# 通信后端，nvidia GPU推薦使用NCCL
dist.barrier() # 等待每個(gè)GPU都運(yùn)行完這個(gè)地方以后再繼續(xù)
在main函數(shù)初始階段，進(jìn)行以下初始化操作。需要注意的是，學(xué)習(xí)率需要根據(jù)使用GPU的張數(shù)增加。在這里使用簡單的倍增方法。
def main(args):
if torch.cuda.is_available() isFalse:
raise EnvironmentError("not find GPU device for training.")
# 初始化各進(jìn)程環(huán)境
init_distributed_mode(args=args)
rank = args.rank
device = torch.device(args.device)
batch_size = args.batch_size
num_classes = args.num_classes
weights_path = args.weights
args.lr *= args.world_size # 學(xué)習(xí)率要根據(jù)并行GPU的數(shù)倍增
實(shí)例化數(shù)據(jù)集可以使用單卡相同的方法，但在sample樣本時(shí)，和單機(jī)不同，需要使用DistributedSampler和BatchSampler。
#給每個(gè)rank對應(yīng)的進(jìn)程分配訓(xùn)練的樣本索引
train_sampler=torch.utils.data.distributed.DistributedSampler(train_data_set)
val_sampler=torch.utils.data.distributed.DistributedSampler(val_data_set)
#將樣本索引每batch_size個(gè)元素組成一個(gè)list
train_batch_sampler=torch.utils.data.BatchSampler(
train_sampler,batch_size,drop_last=True)
DistributedSampler原理如圖所示：假設(shè)當(dāng)前數(shù)據(jù)集有0~10共11個(gè)樣本，使用2塊GPU計(jì)算。首先打亂數(shù)據(jù)順序，然后用 11/2 =6（向上取整），然后6乘以GPU個(gè)數(shù)2 = 12，因?yàn)橹挥?1個(gè)數(shù)據(jù)，所以再把第一個(gè)數(shù)據(jù)（索引為6的數(shù)據(jù)）補(bǔ)到末尾，現(xiàn)在就有12個(gè)數(shù)據(jù)可以均勻分到每塊GPU。然后分配數(shù)據(jù)：間隔將數(shù)據(jù)分配到不同的GPU中。
BatchSampler原理: DistributedSmpler將數(shù)據(jù)分配到兩個(gè)GPU上，以第一個(gè)GPU為例，分到的數(shù)據(jù)是6，9，10，1，8，7，假設(shè)batch_size=2，就按順序把數(shù)據(jù)兩兩一組，在訓(xùn)練時(shí)，每次獲取一個(gè)batch的數(shù)據(jù)，就從組織好的一個(gè)個(gè)batch中取到。注意：只對訓(xùn)練集處理，驗(yàn)證集不使用BatchSampler。
接下來使用定義好的數(shù)據(jù)集和sampler方法加載數(shù)據(jù)：
train_loader = torch.utils.data.DataLoader(train_data_set,
batch_sampler=train_batch_sampler,
pin_memory=True, # 直接加載到顯存中，達(dá)到加速效果
num_workers=nw,
collate_fn=train_data_set.collate_fn)
val_loader = torch.utils.data.DataLoader(val_data_set,
batch_size=batch_size,
sampler=val_sampler,
pin_memory=True,
num_workers=nw,
collate_fn=val_data_set.collate_fn)
如果有預(yù)訓(xùn)練權(quán)重的話，需要保證每塊GPU加載的權(quán)重是一模一樣的。需要在主進(jìn)程保存模型初始化權(quán)重，在不同設(shè)備上載入主進(jìn)程保存的權(quán)重。這樣才能保證每塊GOU上加載的權(quán)重是一致的：
# 實(shí)例化模型
model = resnet34(num_classes=num_classes).to(device)
# 如果存在預(yù)訓(xùn)練權(quán)重則載入
if os.path.exists(weights_path):
weights_dict = torch.load(weights_path, map_location=device)
# 簡單對比每層的權(quán)重參數(shù)個(gè)數(shù)是否一致
load_weights_dict = {k: v for k, v in weights_dict.items()
if model.state_dict()[k].numel() == v.numel()}
model.load_state_dict(load_weights_dict, strict=False)
else:
checkpoint_path = os.path.join(tempfile.gettempdir(), "initial_weights.pt")
# 如果不存在預(yù)訓(xùn)練權(quán)重，需要將第一個(gè)進(jìn)程中的權(quán)重保存，然后其他進(jìn)程載入，保持初始化權(quán)重一致
if rank == 0:
torch.save(model.state_dict(), checkpoint_path)
dist.barrier()
# 這里注意，一定要指定map_location參數(shù)，否則會導(dǎo)致第一塊GPU占用更多資源
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
如果需要凍結(jié)模型權(quán)重，和單GPU基本沒有差別。如果不需要凍結(jié)權(quán)重，可以選擇是否同步BN層。然后再把模型包裝成DDP模型，就可以方便進(jìn)程之間的通信了。多GPU和單GPU的優(yōu)化器設(shè)置沒有差別，這里不再贅述。
# 是否凍結(jié)權(quán)重
if args.freeze_layers:
for name, para in model.named_parameters():
# 除最后的全連接層外，其他權(quán)重全部凍結(jié)
if"fc"notin name:
para.requires_grad_(False)
else:
# 只有訓(xùn)練帶有BN結(jié)構(gòu)的網(wǎng)絡(luò)時(shí)使用SyncBatchNorm采用意義
if args.syncBN:
# 使用SyncBatchNorm后訓(xùn)練會更耗時(shí)
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
# 轉(zhuǎn)為DDP模型
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
# optimizer使用SGD+余弦淬火策略
pg = [p for p in model.parameters() if p.requires_grad]
optimizer = optim.SGD(pg, lr=args.lr, momentum=0.9, weight_decay=0.005)
lf = lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2) * (1 - args.lrf) + args.lrf # cosine
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
與單GPU不同的地方：rain_sampler.set_epoch(epoch)，這行代碼會在每次迭代的時(shí)候獲得一個(gè)不同的生成器，每一輪開始迭代獲取數(shù)據(jù)之前設(shè)置隨機(jī)種子，通過改變傳進(jìn)的epoch參數(shù)改變打亂數(shù)據(jù)順序。通過設(shè)置不同的隨機(jī)種子，可以讓不同GPU每輪拿到的數(shù)據(jù)不同。后面的部分和單GPU相同。
for epoch in range(args.epochs):
train_sampler.set_epoch(epoch)
mean_loss = train_one_epoch(model=model,
optimizer=optimizer,
data_loader=train_loader,
device=device,
epoch=epoch)
scheduler.step()
sum_num = evaluate(model=model,
data_loader=val_loader,
device=device)
acc = sum_num / val_sampler.total_size
我們詳細(xì)看看每個(gè)epoch是訓(xùn)練時(shí)和單GPU訓(xùn)練的差異（上面的train_one_epoch）
def train_one_epoch(model, optimizer, data_loader, device, epoch):
model.train()
loss_function = torch.nn.CrossEntropyLoss()
mean_loss = torch.zeros(1).to(device)
optimizer.zero_grad()
# 在進(jìn)程0中打印訓(xùn)練進(jìn)度
if is_main_process():
data_loader = tqdm(data_loader)
for step, data in enumerate(data_loader):
images, labels = data
pred = model(images.to(device))
loss = loss_function(pred, labels.to(device))
loss.backward()
loss = reduce_value(loss, average=True) # 在單GPU中不起作用，多GPU時(shí)，獲得所有GPU的loss的均值。
mean_loss = (mean_loss * step + loss.detach()) / (step + 1) # update mean losses
# 在進(jìn)程0中打印平均loss
if is_main_process():
data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))
ifnot torch.isfinite(loss):
print('WARNING: non-finite loss, ending training ', loss)
sys.exit(1)
optimizer.step()
optimizer.zero_grad()
# 等待所有進(jìn)程計(jì)算完畢
if device != torch.device("cpu"):
torch.cuda.synchronize(device)
return mean_loss.item()
def reduce_value(value, average=True):
world_size = get_world_size()
if world_size < 2: # 單GPU的情況
return value
with torch.no_grad():
dist.all_reduce(value) # 對不同設(shè)備之間的value求和
if average: # 如果需要求平均，獲得多塊GPU計(jì)算loss的均值
value /= world_size
return value
接下來看一下驗(yàn)證階段的情況，和單GPU最大的額不同之處是預(yù)測正確樣本個(gè)數(shù)的地方。
@torch.no_grad()
def evaluate(model, data_loader, device):
model.eval()
# 用于存儲預(yù)測正確的樣本個(gè)數(shù)，每塊GPU都會計(jì)算自己正確樣本的數(shù)量
sum_num = torch.zeros(1).to(device)
# 在進(jìn)程0中打印驗(yàn)證進(jìn)度
if is_main_process():
data_loader = tqdm(data_loader)
for step, data in enumerate(data_loader):
images, labels = data
pred = model(images.to(device))
pred = torch.max(pred, dim=1)[1]
sum_num += torch.eq(pred, labels.to(device)).sum()
# 等待所有進(jìn)程計(jì)算完畢
if device != torch.device("cpu"):
torch.cuda.synchronize(device)
sum_num = reduce_value(sum_num, average=False) # 預(yù)測正確樣本個(gè)數(shù)
return sum_num.item()
需要注意的是：保存模型的權(quán)重需要在主進(jìn)程中進(jìn)行保存。
if rank == 0:
print("[epoch {}] accuracy: {}".format(epoch, round(acc, 3)))
tags = ["loss", "accuracy", "learning_rate"]
tb_writer.add_scalar(tags[0], mean_loss, epoch)
tb_writer.add_scalar(tags[1], acc, epoch)
tb_writer.add_scalar(tags[2], optimizer.param_groups[0]["lr"], epoch)
torch.save(model.module.state_dict(), "./weights/model-{}.pth".format(epoch))
如果從頭開始訓(xùn)練，主進(jìn)程生成的初始化權(quán)重是以臨時(shí)文件的形式保存，需要訓(xùn)練完后移除掉。最后還需要撤銷進(jìn)程組。
if rank == 0:# 刪除臨時(shí)緩存文件 if os.path.exists(checkpoint_path) is True: os.remove(checkpoint_path) dist.destroy_process_group() # 撤銷進(jìn)程組，釋放資源
鳴謝：本博客內(nèi)容借鑒于up主：霹靂吧啦Wz