實操教程 | GPU多卡并行訓練總結（以pytorch為例）
來源 | https://zhuanlan.zhihu.com/p/402198819
編輯 | 極市平臺
極市導讀
本文的論述分為“為什么要使用多GPU并行訓練”、“常見的多GPU訓練方法”、“誤差梯度如何在不同設備之間通信”、“BN如何在不同設備之間同步”、“兩種GPU訓練方法：DataParallel 和 DistributedDataParallel”、“pytorch中常見的GPU啟動方式”六部分。>>加入極市CV技術交流群，走在計算機視覺的最前沿
為什么要使用多GPU并行訓練
簡單來說，有兩種原因：第一種是模型在一塊GPU上放不下，兩塊或多塊GPU上就能運行完整的模型（如早期的AlexNet）。第二種是多塊GPU并行計算可以達到加速訓練的效果。想要成為“煉丹大師“，多GPU并行訓練是不可或缺的技能。
常見的多GPU訓練方法：
1.模型并行方式：如果模型特別大，GPU顯存不夠，無法將一個顯存放在GPU上，需要把網絡的不同模塊放在不同GPU上，這樣可以訓練比較大的網絡。（下圖左半部分）
2.數據并行方式：將整個模型放在一塊GPU里，再復制到每一塊GPU上，同時進行正向傳播和反向誤差傳播。相當于加大了batch_size。（下圖右半部分）
在pytorch1.7 + cuda10 + TeslaV100的環(huán)境下，使用ResNet34，batch_size=16, SGD對花草數據集訓練的情況如下：使用一塊GPU需要9s一個epoch，使用兩塊GPU是5.5s， 8塊是2s。這里有一個問題，為什么運行時間不是9/8≈1.1s ? 因為使用GPU數量越多，設備之間的通訊會越來越復雜，所以隨著GPU數量的增加，訓練速度的提升也是遞減的。
誤差梯度如何在不同設備之間通信？
在每個GPU訓練step結束后，將每塊GPU的損失梯度求平均，而不是每塊GPU各計算各的。
BN如何在不同設備之間同步？
假設batch_size=2，每個GPU計算的均值和方差都針對這兩個樣本而言的。而BN的特性是：batch_size越大，均值和方差越接近與整個數據集的均值和方差，效果越好。使用多塊GPU時，會計算每個BN層在所有設備上輸入的均值和方差。如果GPU1和GPU2都分別得到兩個特征層，那么兩塊GPU一共計算4個特征層的均值和方差，可以認為batch_size=4。注意：如果不用同步BN，而是每個設備計算自己的批次數據的均值方差，效果與單GPU一致，僅僅能提升訓練速度；如果使用同步BN，效果會有一定提升，但是會損失一部分并行速度。
下圖為單GPU、以及是否使用同步BN訓練的三種情況，可以看到使用同步BN（橙線）比不使用同步BN（藍線）總體效果要好一些，不過訓練時間也會更長。使用單GPU（黑線）和不使用同步BN的效果是差不多的。
兩種GPU訓練方法：DataParallel 和 DistributedDataParallel：
DataParallel是單進程多線程的，僅僅能工作在單機中。而DistributedDataParallel是多進程的，可以工作在單機或多機器中。 DataParallel通常會慢于DistributedDataParallel。所以目前主流的方法是DistributedDataParallel。
pytorch中常見的GPU啟動方式：
注：distributed.launch方法如果開始訓練后，手動終止程序，最好先看下顯存占用情況，有小概率進程沒kill的情況，會占用一部分GPU顯存資源。
下面以分類問題為基準，詳細介紹使用DistributedDataParallel時的過程:
首先要初始化各進程環(huán)境：
def init_distributed_mode(args):
# 如果是多機多卡的機器，WORLD_SIZE代表使用的機器數，RANK對應第幾臺機器
# 如果是單機多卡的機器，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代表某個機器上第幾塊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) # 對當前進程指定使用的GPU
args.dist_backend = 'nccl' # 通信后端，nvidia GPU推薦使用NCCL
dist.barrier() # 等待每個GPU都運行完這個地方以后再繼續(xù)
在main函數初始階段，進行以下初始化操作。需要注意的是，學習率需要根據使用GPU的張數增加。在這里使用簡單的倍增方法。
def main(args):
if torch.cuda.is_available() is False:
raise EnvironmentError("not find GPU device for training.")
# 初始化各進程環(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 # 學習率要根據并行GPU的數倍增
實例化數據集可以使用單卡相同的方法，但在sample樣本時，和單機不同，需要使用DistributedSampler和BatchSampler。
#給每個rank對應的進程分配訓練的樣本索引
train_sampler=torch.utils.data.distributed.DistributedSampler(train_data_set)
val_sampler=torch.utils.data.distributed.DistributedSampler(val_data_set)
#將樣本索引每batch_size個元素組成一個list
train_batch_sampler=torch.utils.data.BatchSampler(
train_sampler,batch_size,drop_last=True)
DistributedSampler原理如圖所示：假設當前數據集有0~10共11個樣本，使用2塊GPU計算。首先打亂數據順序，然后用 11/2 =6（向上取整），然后6乘以GPU個數2 = 12，因為只有11個數據，所以再把第一個數據（索引為6的數據）補到末尾，現在就有12個數據可以均勻分到每塊GPU。然后分配數據：間隔將數據分配到不同的GPU中。
BatchSampler原理: DistributedSmpler將數據分配到兩個GPU上，以第一個GPU為例，分到的數據是6，9，10，1，8，7，假設batch_size=2，就按順序把數據兩兩一組，在訓練時，每次獲取一個batch的數據，就從組織好的一個個batch中取到。注意：只對訓練集處理，驗證集不使用BatchSampler。
接下來使用定義好的數據集和sampler方法加載數據：
train_loader = torch.utils.data.DataLoader(train_data_set,
batch_sampler=train_batch_sampler,
pin_memory=True, # 直接加載到顯存中，達到加速效果
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)
如果有預訓練權重的話，需要保證每塊GPU加載的權重是一模一樣的。需要在主進程保存模型初始化權重，在不同設備上載入主進程保存的權重。這樣才能保證每塊GOU上加載的權重是一致的：
# 實例化模型
model = resnet34(num_classes=num_classes).to(device)
# 如果存在預訓練權重則載入
if os.path.exists(weights_path):
weights_dict = torch.load(weights_path, map_location=device)
# 簡單對比每層的權重參數個數是否一致
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")
# 如果不存在預訓練權重，需要將第一個進程中的權重保存，然后其他進程載入，保持初始化權重一致
if rank == 0:
torch.save(model.state_dict(), checkpoint_path)
dist.barrier()
# 這里注意，一定要指定map_location參數，否則會導致第一塊GPU占用更多資源
model.load_state_dict(torch.load(checkpoint_path, map_location=device))
如果需要凍結模型權重，和單GPU基本沒有差別。如果不需要凍結權重，可以選擇是否同步BN層。然后再把模型包裝成DDP模型，就可以方便進程之間的通信了。多GPU和單GPU的優(yōu)化器設置沒有差別，這里不再贅述。
# 是否凍結權重
if args.freeze_layers:
for name, para in model.named_parameters():
# 除最后的全連接層外，其他權重全部凍結
if "fc" not in name:
para.requires_grad_(False)
else:
# 只有訓練帶有BN結構的網絡時使用SyncBatchNorm采用意義
if args.syncBN:
# 使用SyncBatchNorm后訓練會更耗時
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
# 轉為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)，這行代碼會在每次迭代的時候獲得一個不同的生成器，每一輪開始迭代獲取數據之前設置隨機種子，通過改變傳進的epoch參數改變打亂數據順序。通過設置不同的隨機種子，可以讓不同GPU每輪拿到的數據不同。后面的部分和單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
我們詳細看看每個epoch是訓練時和單GPU訓練的差異（上面的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()
# 在進程0中打印訓練進度
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時，獲得所有GPU的loss的均值。
mean_loss = (mean_loss * step + loss.detach()) / (step + 1) # update mean losses
# 在進程0中打印平均loss
if is_main_process():
data_loader.desc = "[epoch {}] mean loss {}".format(epoch, round(mean_loss.item(), 3))
if not torch.isfinite(loss):
print('WARNING: non-finite loss, ending training ', loss)
sys.exit(1)
optimizer.step()
optimizer.zero_grad()
# 等待所有進程計算完畢
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) # 對不同設備之間的value求和
if average: # 如果需要求平均，獲得多塊GPU計算loss的均值
value /= world_size
return value
接下來看一下驗證階段的情況，和單GPU最大的額不同之處是預測正確樣本個數的地方。
@torch.no_grad()
def evaluate(model, data_loader, device):
model.eval()
# 用于存儲預測正確的樣本個數，每塊GPU都會計算自己正確樣本的數量
sum_num = torch.zeros(1).to(device)
# 在進程0中打印驗證進度
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()
# 等待所有進程計算完畢
if device != torch.device("cpu"):
torch.cuda.synchronize(device)
sum_num = reduce_value(sum_num, average=False) # 預測正確樣本個數
return sum_num.item()
需要注意的是：保存模型的權重需要在主進程中進行保存。
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))
如果從頭開始訓練，主進程生成的初始化權重是以臨時文件的形式保存，需要訓練完后移除掉。最后還需要撤銷進程組。
if rank == 0:# 刪除臨時緩存文件 if os.path.exists(checkpoint_path) is True: os.remove(checkpoint_path) dist.destroy_process_group() # 撤銷進程組，釋放資源
鳴謝：本博客內容借鑒于up主：霹靂吧啦Wz
如果覺得有用，就請分享到朋友圈吧！
公眾號后臺回復“85”獲取ICCV2021 oral直播分享PPT下載～
# CV技術社群邀請函 #
備注：姓名-學校/公司-研究方向-城市（如：小極-北大-目標檢測-深圳）
即可申請加入極市目標檢測/圖像分割/工業(yè)檢測/人臉/醫(yī)學影像/3D/SLAM/自動駕駛/超分辨率/姿態(tài)估計/ReID/GAN/圖像增強/OCR/視頻理解等技術交流群
每月大咖直播分享、真實項目需求對接、求職內推、算法競賽、干貨資訊匯總、與 10000+來自港科大、北大、清華、中科院、CMU、騰訊、百度等名校名企視覺開發(fā)者互動交流~
