PyTorch (高級(jí)）帶有 Amazon AWS 的 PyTorch 1.0 分布式訓(xùn)練師

2020-09-10 10:31 更新
原文： https://pytorch.org/tutorials/beginner/aws_distributed_training_tutorial.html

作者： Nathan Inkawhich

由編輯：滕力

在本教程中，我們將展示如何在兩個(gè)多 GPU Amazon AWS 節(jié)點(diǎn)之間設(shè)置，編碼和運(yùn)行 PyTorch 1.0 分布式訓(xùn)練師。 我們將首先描述 AWS 設(shè)置，然后是 PyTorch 環(huán)境配置，最后是分布式訓(xùn)練師的代碼。 希望您會(huì)發(fā)現(xiàn)將當(dāng)前的訓(xùn)練代碼擴(kuò)展到分布式應(yīng)用程序?qū)嶋H上只需很少的代碼更改，并且大部分工作都在一次性環(huán)境設(shè)置中。

Amazon AWS 安裝程序

在本教程中，我們將在兩個(gè)多 GPU 節(jié)點(diǎn)之間進(jìn)行分布式訓(xùn)練。 在本節(jié)中，我們將首先介紹如何創(chuàng)建節(jié)點(diǎn)，然后介紹如何設(shè)置安全組，以便節(jié)點(diǎn)之間可以相互通信。

創(chuàng)建節(jié)點(diǎn)

在 Amazon AWS 中，有七個(gè)創(chuàng)建實(shí)例的步驟。 首先，登錄并選擇 Launch Instance 。

步驟 1：選擇一個(gè)亞馬遜機(jī)器映像(AMI）-在這里，我們將選擇Deep Learning AMI (Ubuntu) Version 14.0。 如前所述，此實(shí)例隨附了許多最受歡迎的深度學(xué)習(xí)框架，并且已預(yù)先配置了 CUDA，cuDNN 和 NCCL。 這是本教程的一個(gè)很好的起點(diǎn)。

步驟 2：選擇實(shí)例類(lèi)型-現(xiàn)在，選擇名為p2.8xlarge的 GPU 計(jì)算單元。 請(qǐng)注意，每個(gè)實(shí)例的成本都不同，但是此實(shí)例每個(gè)節(jié)點(diǎn)提供 8 個(gè) NVIDIA Tesla K80 GPU，并為多 GPU 分布式訓(xùn)練提供了良好的體系結(jié)構(gòu)。

步驟 3：配置實(shí)例詳細(xì)信息-此處唯一要更改的設(shè)置是將實(shí)例數(shù)增至 2。默認(rèn)情況下，所有其他配置都可以保留。

步驟 4：添加存儲(chǔ)-注意，默認(rèn)情況下，這些節(jié)點(diǎn)沒(méi)有很多存儲(chǔ)(只有 75 GB）。 對(duì)于本教程，由于我們僅使用 STL-10 數(shù)據(jù)集，因此有足夠的存儲(chǔ)空間。 但是，如果要在較大的數(shù)據(jù)集(如 ImageNet）上進(jìn)行訓(xùn)練，則必須添加更多的存儲(chǔ)空間以適合數(shù)據(jù)集和要保存的任何訓(xùn)練后的模型。

步驟 5：添加標(biāo)簽-此處無(wú)需執(zhí)行任何操作，只需繼續(xù)。

步驟 6：配置安全組-這是配置過(guò)程中的關(guān)鍵步驟。 默認(rèn)情況下，同一安全組中的兩個(gè)節(jié)點(diǎn)將無(wú)法在分布式訓(xùn)練設(shè)置中進(jìn)行通信。 在這里，我們要為要加入的兩個(gè)節(jié)點(diǎn)創(chuàng)建一個(gè)新的安全組。但是，我們無(wú)法在此步驟中完成配置。 現(xiàn)在，只需記住您的新安全組名稱(chēng)(例如 launch-wizard-12），然后繼續(xù)執(zhí)行步驟 7。

步驟 7：查看實(shí)例啟動(dòng)-在這里，查看實(shí)例然后啟動(dòng)它。 默認(rèn)情況下，這將自動(dòng)開(kāi)始初始化兩個(gè)實(shí)例。 您可以從儀表板監(jiān)視初始化進(jìn)度。

配置安全組

回想一下，我們?cè)趧?chuàng)建實(shí)例時(shí)無(wú)法正確配置安全組。 啟動(dòng)實(shí)例后，在 EC2 儀表板中選擇網(wǎng)絡(luò)&安全>安全組選項(xiàng)卡。 這將顯示您有權(quán)訪問(wèn)的安全組的列表。 選擇您在步驟 6 中創(chuàng)建的新安全組(即 launch-wizard-12），該安全組將彈出名為 Description，Inbound，Outbound 和 Tags 的標(biāo)簽。 首先，選擇入站選項(xiàng)卡，然后選擇編輯，添加一個(gè)規(guī)則，以允許 launch-wizard-12 安全組中“源”的“所有流量”。 然后選擇出站選項(xiàng)卡，并執(zhí)行完全相同的操作。 現(xiàn)在，我們已經(jīng)有效地允許了 launch-wizard-12 安全組中節(jié)點(diǎn)之間所有類(lèi)型的所有入站和出站流量。

必要信息

在繼續(xù)之前，我們必須找到并記住兩個(gè)節(jié)點(diǎn)的 IP 地址。 在 EC2 儀表板中找到正在運(yùn)行的實(shí)例。 對(duì)于這兩種情況，記下 IPv4 公用 IP 和專(zhuān)用 IP 。 在本文檔的其余部分，我們將它們稱(chēng)為 node0-publicIP ， node0-privateIP ， node1-publicIP 和 node1- privateIP 。 公用 IP 是我們將用于 SSH 的地址，而專(zhuān)用 IP 將用于節(jié)點(diǎn)間的通信。

環(huán)境設(shè)定

下一個(gè)關(guān)鍵步驟是每個(gè)節(jié)點(diǎn)的設(shè)置。 不幸的是，我們不能同時(shí)配置兩個(gè)節(jié)點(diǎn)，因此必須在每個(gè)節(jié)點(diǎn)上分別完成此過(guò)程。 但是，這是一次設(shè)置，因此一旦正確配置了節(jié)點(diǎn)，就無(wú)需為將來(lái)的分布式訓(xùn)練項(xiàng)目重新配置。

一旦登錄到節(jié)點(diǎn)，第一步就是使用 python 3.6 和 numpy 創(chuàng)建一個(gè)新的 conda 環(huán)境。 創(chuàng)建后，激活環(huán)境。

  1. $ conda create -n nightly_pt python=3.6 numpy
  2. $ source activate nightly_pt

接下來(lái)，我們將在 conda 環(huán)境中安裝啟用了 Cuda 9.0 的 PyTorch 的每晚構(gòu)建。

  1. $ pip install torch_nightly -f https://download.pytorch.org/whl/nightly/cu90/torch_nightly.html

我們還必須安裝 torchvision，以便可以使用 torchvision 模型和數(shù)據(jù)集。 目前，我們必須從源代碼構(gòu)建 torchvision，因?yàn)槟J(rèn)情況下，pip 安裝將在我們剛剛安裝的每晚構(gòu)建的基礎(chǔ)上安裝舊版本的 PyTorch。

  1. $ cd
  2. $ git clone https://github.com/pytorch/vision.git
  3. $ cd vision
  4. $ python setup.py install

最后，非常重要的步驟是設(shè)置 NCCL 套接字的網(wǎng)絡(luò)接口名稱(chēng)。 這是通過(guò)環(huán)境變量NCCL_SOCKET_IFNAME設(shè)置的。 要獲得正確的名稱(chēng)，請(qǐng)?jiān)诠?jié)點(diǎn)上運(yùn)行ifconfig命令，然后查看與該節(jié)點(diǎn)的 privateIP 相對(duì)應(yīng)的接口名稱(chēng)(例如 ens3）。 然后將環(huán)境變量設(shè)置為

  1. $ export NCCL_SOCKET_IFNAME=ens3

請(qǐng)記住，在兩個(gè)節(jié)點(diǎn)上都執(zhí)行此操作。 您也可以考慮將 NCCLSOCKET_IFNAME 設(shè)置添加到 .bashrc_ 中。 一個(gè)重要的觀察結(jié)果是我們沒(méi)有在節(jié)點(diǎn)之間設(shè)置共享文件系統(tǒng)。 因此，每個(gè)節(jié)點(diǎn)都必須具有代碼的副本和數(shù)據(jù)集的副本。 有關(guān)在節(jié)點(diǎn)之間設(shè)置共享網(wǎng)絡(luò)文件系統(tǒng)的更多信息，請(qǐng)參見(jiàn)。

分布式訓(xùn)練守則

通過(guò)運(yùn)行實(shí)例和環(huán)境設(shè)置，我們現(xiàn)在可以進(jìn)入訓(xùn)練代碼。 這里的大多數(shù)代碼均來(lái)自 PyTorch ImageNet 示例，該示例也支持分布式訓(xùn)練。 該代碼為定制訓(xùn)練師提供了一個(gè)很好的起點(diǎn)，因?yàn)樗哂性S多樣板訓(xùn)練循環(huán)，驗(yàn)證循環(huán)和準(zhǔn)確性跟蹤功能。 但是，您會(huì)注意到，為簡(jiǎn)單起見(jiàn)，已刪除了參數(shù)解析和其他非必要功能。

在此示例中，我們將使用 torchvision.models.resnet18 模型，并將其訓(xùn)練在 torchvision.datasets.STL10 數(shù)據(jù)集上。 為了適應(yīng) STL-10 與 Resnet18 的尺寸不匹配，我們將通過(guò)變換將每個(gè)圖像的大小調(diào)整為 224x224。 注意，模型和數(shù)據(jù)集的選擇與分布式訓(xùn)練代碼正交，您可以使用所需的任何數(shù)據(jù)集和模型，并且過(guò)程相同。 首先處理導(dǎo)入，然后討論一些輔助函數(shù)。 然后，我們將定義訓(xùn)練和測(cè)試功能，這些功能很大程度上取自 ImageNet 示例。 最后，我們將構(gòu)建處理分布式訓(xùn)練設(shè)置的代碼的主要部分。 最后，我們將討論如何實(shí)際運(yùn)行代碼。

進(jìn)口貨

這里重要的分布式訓(xùn)練專(zhuān)用導(dǎo)入是 torch.nn.parallel ， torch.distributed ，  torch.utils.data.distributed 和torch.并行處理。 將并行處理啟動(dòng)方法設(shè)置為生成或 forkserver (僅在 Python 3 中受支持），這一點(diǎn)也很重要，因?yàn)槟J(rèn)值為 fork ，這在發(fā)生以下情況時(shí)可能導(dǎo)致死鎖 使用多個(gè)工作進(jìn)程進(jìn)行數(shù)據(jù)加載。

  1. import time
  2. import sys
  3. import torch
  6. import torch.nn as nn
  7. import torch.nn.parallel
  8. import torch.distributed as dist
  9. import torch.optim
  10. import torch.utils.data
  11. import torch.utils.data.distributed
  12. import torchvision.transforms as transforms
  13. import torchvision.datasets as datasets
  14. import torchvision.models as models
  17. from torch.multiprocessing import Pool, Process

輔助功能

我們還必須定義一些輔助函數(shù)和類(lèi)，以使訓(xùn)練更加容易。 AverageMeter類(lèi)跟蹤訓(xùn)練統(tǒng)計(jì)信息，例如準(zhǔn)確性和迭代計(jì)數(shù)。 accuracy函數(shù)計(jì)算并返回模型的 top-k 精度，因此我們可以跟蹤學(xué)習(xí)進(jìn)度。 兩者都是為了訓(xùn)練方便而提供的，但都不是專(zhuān)門(mén)針對(duì)分布式訓(xùn)練的。

  1. class AverageMeter(object):
  2.     """Computes and stores the average and current value"""
  3.     def __init__(self):
  4.         self.reset()
  7.     def reset(self):
  8.         self.val = 0
  9.         self.avg = 0
  10.         self.sum = 0
  11.         self.count = 0
  14.     def update(self, val, n=1):
  15.         self.val = val
  16.         self.sum += val * n
  17.         self.count += n
  18.         self.avg = self.sum / self.count
  21. def accuracy(output, target, topk=(1,)):
  22.     """Computes the precision@k for the specified values of k"""
  23.     with torch.no_grad():
  24.         maxk = max(topk)
  25.         batch_size = target.size(0)
  28.         _, pred = output.topk(maxk, 1, True, True)
  29.         pred = pred.t()
  30.         correct = pred.eq(target.view(1, -1).expand_as(pred))
  33.         res = []
  34.         for k in topk:
  35.             correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
  36.             res.append(correct_k.mul_(100.0 / batch_size))
  37.         return res

火車(chē)功能

為了簡(jiǎn)化主循環(huán)，最好將訓(xùn)練紀(jì)元步驟分離為一個(gè)稱(chēng)為train的函數(shù)。 此函數(shù)為 train_loader 的一個(gè)時(shí)期訓(xùn)練輸入模型。 此功能中唯一的分布式訓(xùn)練工件是在正向傳遞之前將數(shù)據(jù)和標(biāo)簽張量的 non_blocking 屬性設(shè)置為True。 這允許數(shù)據(jù)的異步 GPU 副本，意味著傳輸可以與計(jì)算重疊。 此功能還會(huì)沿途輸出訓(xùn)練統(tǒng)計(jì)信息，以便我們可以跟蹤整個(gè)時(shí)期的進(jìn)度。

在此定義的另一個(gè)功能是adjust_learning_rate，它以固定的時(shí)間表衰減初始學(xué)習(xí)率。 這是另一個(gè)樣板訓(xùn)練器功能，可用于訓(xùn)練準(zhǔn)確的模型。

  1. def train(train_loader, model, criterion, optimizer, epoch):
  4.     batch_time = AverageMeter()
  5.     data_time = AverageMeter()
  6.     losses = AverageMeter()
  7.     top1 = AverageMeter()
  8.     top5 = AverageMeter()
  11.     # switch to train mode
  12.     model.train()
  15.     end = time.time()
  16.     for i, (input, target) in enumerate(train_loader):
  19.         # measure data loading time
  20.         data_time.update(time.time() - end)
  23.         # Create non_blocking tensors for distributed training
  24.         input = input.cuda(non_blocking=True)
  25.         target = target.cuda(non_blocking=True)
  28.         # compute output
  29.         output = model(input)
  30.         loss = criterion(output, target)
  33.         # measure accuracy and record loss
  34.         prec1, prec5 = accuracy(output, target, topk=(1, 5))
  35.         losses.update(loss.item(), input.size(0))
  36.         top1.update(prec1[0], input.size(0))
  37.         top5.update(prec5[0], input.size(0))
  40.         # compute gradients in a backward pass
  41.         optimizer.zero_grad()
  42.         loss.backward()
  45.         # Call step of optimizer to update model params
  46.         optimizer.step()
  49.         # measure elapsed time
  50.         batch_time.update(time.time() - end)
  51.         end = time.time()
  54.         if i % 10 == 0:
  55.             print('Epoch: [{0}][{1}/{2}]\t'
  56.                   'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
  57.                   'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
  58.                   'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
  59.                   'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
  60.                   'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
  61.                    epoch, i, len(train_loader), batch_time=batch_time,
  62.                    data_time=data_time, loss=losses, top1=top1, top5=top5))
  65. def adjust_learning_rate(initial_lr, optimizer, epoch):
  66.     """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
  67.     lr = initial_lr * (0.1 ** (epoch // 30))
  68.     for param_group in optimizer.param_groups:
  69.         param_group['lr'] = lr

驗(yàn)證功能

為了跟蹤泛化性能并進(jìn)一步簡(jiǎn)化主循環(huán)，我們還可以將驗(yàn)證步驟提取到一個(gè)名為validate的函數(shù)中。 此函數(shù)在輸入驗(yàn)證數(shù)據(jù)加載器上運(yùn)行輸入模型的完整驗(yàn)證步驟，并在驗(yàn)證集上返回模型的 top-1 準(zhǔn)確性。 再次，您會(huì)注意到這里唯一的分布式訓(xùn)練功能是在將訓(xùn)練數(shù)據(jù)和標(biāo)簽傳遞到模型之前為它們?cè)O(shè)置non_blocking=True。

  1. def validate(val_loader, model, criterion):
  4.     batch_time = AverageMeter()
  5.     losses = AverageMeter()
  6.     top1 = AverageMeter()
  7.     top5 = AverageMeter()
  10.     # switch to evaluate mode
  11.     model.eval()
  14.     with torch.no_grad():
  15.         end = time.time()
  16.         for i, (input, target) in enumerate(val_loader):
  19.             input = input.cuda(non_blocking=True)
  20.             target = target.cuda(non_blocking=True)
  23.             # compute output
  24.             output = model(input)
  25.             loss = criterion(output, target)
  28.             # measure accuracy and record loss
  29.             prec1, prec5 = accuracy(output, target, topk=(1, 5))
  30.             losses.update(loss.item(), input.size(0))
  31.             top1.update(prec1[0], input.size(0))
  32.             top5.update(prec5[0], input.size(0))
  35.             # measure elapsed time
  36.             batch_time.update(time.time() - end)
  37.             end = time.time()
  40.             if i % 100 == 0:
  41.                 print('Test: [{0}/{1}]\t'
  42.                       'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
  43.                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
  44.                       'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
  45.                       'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
  46.                        i, len(val_loader), batch_time=batch_time, loss=losses,
  47.                        top1=top1, top5=top5))
  50.         print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
  51.               .format(top1=top1, top5=top5))
  54.     return top1.avg

輸入項(xiàng)

借助輔助功能，現(xiàn)在我們到達(dá)了有趣的部分。 在這里，我們將定義運(yùn)行的輸入。 一些輸入是標(biāo)準(zhǔn)模型訓(xùn)練輸入，例如批次大小和訓(xùn)練時(shí)期數(shù)，而某些則特定于我們的分布式訓(xùn)練任務(wù)。 所需的輸入是：

  • batch_size -分布式訓(xùn)練組中每個(gè)進(jìn)程的的批處理大小。 整個(gè)分布式模型的總批次大小為 batch_size * world_size
  • worker -每個(gè)進(jìn)程中與數(shù)據(jù)加載器一起使用的工作進(jìn)程數(shù)
  • num_epochs -要訓(xùn)練的時(shí)期總數(shù)
  • starting_lr -訓(xùn)練的開(kāi)始學(xué)習(xí)率
  • world_size -分布式訓(xùn)練環(huán)境中的進(jìn)程數(shù)
  • dist_backend -用于分布式訓(xùn)練通信(即 NCCL，Glo，MPI 等）的后端。 在本教程中，由于我們使用了多個(gè)多 GPU 節(jié)點(diǎn)，因此建議使用 NCCL。
  • dist_url -用于指定進(jìn)程組的初始化方法的 URL。 它可能包含 rank0 進(jìn)程的 IP 地址和端口，或者是共享文件系統(tǒng)上不存在的文件。 在這里，由于我們沒(méi)有共享文件系統(tǒng)，因此它將合并 node0-privateIP 和 node0 上要使用的端口。
  1. print("Collect Inputs...")
  4. ## Batch Size for training and testing
  5. batch_size = 32
  8. ## Number of additional worker processes for dataloading
  9. workers = 2
  12. ## Number of epochs to train for
  13. num_epochs = 2
  16. ## Starting Learning Rate
  17. starting_lr = 0.1
  20. ## Number of distributed processes
  21. world_size = 4
  24. ## Distributed backend type
  25. dist_backend = 'nccl'
  28. ## Url used to setup distributed training
  29. dist_url = "tcp://172.31.22.234:23456"

初始化流程組

在 PyTorch 中，分布式訓(xùn)練最重要的部分之一就是正確設(shè)置過(guò)程組，這是初始化torch.distributed包時(shí)首先執(zhí)行的步驟。 為此，我們將使用torch.distributed.init_process_group函數(shù)，該函數(shù)需要多個(gè)輸入。 首先，后端輸入指定要使用的后端(即 NCCL，Gloo，MPI 等）。 init_method 輸入，該 URL 是包含 rank0 計(jì)算機(jī)的地址和端口的 url，或者是共享文件系統(tǒng)上不存在的文件的路徑。 注意，要使用文件 initmethod，所有計(jì)算機(jī)都必須有權(quán)訪問(wèn)該文件，對(duì)于 url 方法而言，所有計(jì)算機(jī)都必須能夠在網(wǎng)絡(luò)上進(jìn)行通信，因此請(qǐng)確保配置所有防火墻和網(wǎng)絡(luò)設(shè)置以使其適應(yīng)。 _init_process_group 函數(shù)還采用等級(jí)和 world_size 自變量，它們分別指定運(yùn)行時(shí)此進(jìn)程的等級(jí)和集合中的進(jìn)程數(shù)。 init_method 輸入也可以是“ env：//”。 在這種情況下，將分別從以下兩個(gè)環(huán)境變量中讀取 rank0 機(jī)器的地址和端口：MASTERADDR，MASTER_PORT。 如果在 _init_process_group 函數(shù)中未指定等級(jí)和 world_size 參數(shù)，則也可以分別從以下兩個(gè)環(huán)境變量中讀取它們：RANK，WORLD_SIZE。

另一個(gè)重要步驟(尤其是當(dāng)每個(gè)節(jié)點(diǎn)具有多個(gè) GPU 時(shí)）是設(shè)置此過(guò)程的 local_rank 。 例如，如果您有兩個(gè)節(jié)點(diǎn)，每個(gè)節(jié)點(diǎn)有 8 個(gè) GPU，并且您希望對(duì)其全部進(jìn)行訓(xùn)練，則 ，每個(gè)節(jié)點(diǎn)將具有本地等級(jí) 0-7 的進(jìn)程。 此 local_rank 用于設(shè)置進(jìn)程的設(shè)備(即使用哪個(gè) GPU），后來(lái)在創(chuàng)建分布式數(shù)據(jù)并行模型時(shí)用于設(shè)置設(shè)備。 還建議在此假設(shè)環(huán)境中使用 NCCL 后端，因?yàn)閷?duì)于多 GPU 節(jié)點(diǎn)，首選 NCCL。

  1. print("Initialize Process Group...")
  2. ## Initialize Process Group
  3. ## v1 - init with url
  4. dist.init_process_group(backend=dist_backend, init_method=dist_url, rank=int(sys.argv[1]), world_size=world_size)
  5. ## v2 - init with file
  6. ## dist.init_process_group(backend="nccl", init_method="file:///home/ubuntu/pt-distributed-tutorial/trainfile", rank=int(sys.argv[1]), world_size=world_size)
  7. ## v3 - init with environment variables
  8. ## dist.init_process_group(backend="nccl", init_method="env://", rank=int(sys.argv[1]), world_size=world_size)
  11. ## Establish Local Rank and set device on this node
  12. local_rank = int(sys.argv[2])
  13. dp_device_ids = [local_rank]
  14. torch.cuda.set_device(local_rank)

初始化模型

下一步是初始化要訓(xùn)練的模型。 在這里，我們將使用torchvision.models中的 resnet18 模型，但可以使用任何模型。 首先，我們初始化模型并將其放置在 GPU 內(nèi)存中。 接下來(lái)，我們制作模型DistributedDataParallel，該模型處理與模型之間的數(shù)據(jù)分配，這對(duì)于分布式訓(xùn)練至關(guān)重要。 DistributedDataParallel模塊還可以處理世界范圍內(nèi)的梯度平均，因此我們不必在訓(xùn)練步驟中明確地對(duì)梯度進(jìn)行平均。

重要的是要注意，這是一個(gè)阻塞函數(shù)，這意味著程序執(zhí)行將在此函數(shù)等待，直到 world_size 進(jìn)程加入進(jìn)程組為止。 另外，請(qǐng)注意，我們將設(shè)備 ID 列表作為參數(shù)傳遞，其中包含我們正在使用的本地排名(即 GPU）。 最后，我們指定損失函數(shù)和優(yōu)化器來(lái)訓(xùn)練模型。

  1. print("Initialize Model...")
  2. ## Construct Model
  3. model = models.resnet18(pretrained=False).cuda()
  4. ## Make model DistributedDataParallel
  5. model = torch.nn.parallel.DistributedDataParallel(model, device_ids=dp_device_ids, output_device=local_rank)
  8. ## define loss function (criterion) and optimizer
  9. criterion = nn.CrossEntropyLoss().cuda()
  10. optimizer = torch.optim.SGD(model.parameters(), starting_lr, momentum=0.9, weight_decay=1e-4)

初始化數(shù)據(jù)加載器

準(zhǔn)備訓(xùn)練的最后一步是指定要使用的數(shù)據(jù)集。 這里，我們使用 torchvision.datasets.STL10 中的 STL-10 數(shù)據(jù)集。 STL10 數(shù)據(jù)集是 96x96px 彩色圖像的 10 類(lèi)數(shù)據(jù)集。 為了與我們的模型一起使用，我們?cè)谵D(zhuǎn)換中將圖像的大小調(diào)整為 224x224px。 本節(jié)中的一項(xiàng)分布式訓(xùn)練特定項(xiàng)目是將DistributedSampler用于訓(xùn)練集，該訓(xùn)練集旨在與DistributedDataParallel模型結(jié)合使用。 該對(duì)象處理整個(gè)分布式環(huán)境中數(shù)據(jù)集的分區(qū)，因此并非所有模型都在同一數(shù)據(jù)子集上進(jìn)行訓(xùn)練，這會(huì)適得其反。 最后，我們創(chuàng)建DataLoader，負(fù)責(zé)將數(shù)據(jù)饋送到流程中。

如果節(jié)點(diǎn)不存在，STL-10 數(shù)據(jù)集將自動(dòng)在節(jié)點(diǎn)上下載。 如果您希望使用自己的數(shù)據(jù)集，則應(yīng)下載數(shù)據(jù)，編寫(xiě)自己的數(shù)據(jù)集處理程序，并在此處為數(shù)據(jù)集構(gòu)造一個(gè)數(shù)據(jù)加載器。

  1. print("Initialize Dataloaders...")
  2. ## Define the transform for the data. Notice, we must resize to 224x224 with this dataset and model.
  3. transform = transforms.Compose(
  4.     [transforms.Resize(224),
  5.      transforms.ToTensor(),
  6.      transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
  9. ## Initialize Datasets. STL10 will automatically download if not present
  10. trainset = datasets.STL10(root='./data', split='train', download=True, transform=transform)
  11. valset = datasets.STL10(root='./data', split='test', download=True, transform=transform)
  14. ## Create DistributedSampler to handle distributing the dataset across nodes when training
  15. ## This can only be called after torch.distributed.init_process_group is called
  16. train_sampler = torch.utils.data.distributed.DistributedSampler(trainset)
  19. ## Create the Dataloaders to feed data to the training and validation steps
  20. train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=(train_sampler is None), num_workers=workers, pin_memory=False, sampler=train_sampler)
  21. val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False, num_workers=workers, pin_memory=False)

訓(xùn)練循環(huán)

最后一步是定義訓(xùn)練循環(huán)。 我們已經(jīng)完成了設(shè)置分布式訓(xùn)練的大部分工作，因此這不是特定于分布式訓(xùn)練的。 唯一的細(xì)節(jié)是在DistributedSampler中設(shè)置當(dāng)前紀(jì)元計(jì)數(shù)，因?yàn)椴蓸悠鲿?huì)根據(jù)紀(jì)元確定性地隨機(jī)整理進(jìn)入每個(gè)進(jìn)程的數(shù)據(jù)。 更新采樣器后，循環(huán)將運(yùn)行一個(gè)完整的訓(xùn)練時(shí)期，運(yùn)行一個(gè)完整的驗(yàn)證步驟，然后將當(dāng)前模型的性能與迄今為止性能最佳的模型進(jìn)行比較。 訓(xùn)練完 num_epochs 之后，循環(huán)退出，并且教程已完成。 注意，由于這是一項(xiàng)練習(xí)，因此我們沒(méi)有保存模型，但是可能希望跟蹤性能最佳的模型，然后在訓(xùn)練結(jié)束時(shí)保存模型(請(qǐng)參見(jiàn) ，在此處）。

  1. best_prec1 = 0
  4. for epoch in range(num_epochs):
  5.     # Set epoch count for DistributedSampler
  6.     train_sampler.set_epoch(epoch)
  9.     # Adjust learning rate according to schedule
  10.     adjust_learning_rate(starting_lr, optimizer, epoch)
  13.     # train for one epoch
  14.     print("\nBegin Training Epoch {}".format(epoch+1))
  15.     train(train_loader, model, criterion, optimizer, epoch)
  18.     # evaluate on validation set
  19.     print("Begin Validation @ Epoch {}".format(epoch+1))
  20.     prec1 = validate(val_loader, model, criterion)
  23.     # remember best prec@1 and save checkpoint if desired
  24.     # is_best = prec1 > best_prec1
  25.     best_prec1 = max(prec1, best_prec1)
  28.     print("Epoch Summary: ")
  29.     print("\tEpoch Accuracy: {}".format(prec1))
  30.     print("\tBest Accuracy: {}".format(best_prec1))

運(yùn)行代碼

與大多數(shù)其他 PyTorch 教程不同，此代碼可能無(wú)法直接在筆記本中運(yùn)行。 要運(yùn)行，請(qǐng)下載此文件的.py 版本(或使用此進(jìn)行轉(zhuǎn)換），然后將副本上載到兩個(gè)節(jié)點(diǎn)。 精明的讀者會(huì)注意到，我們對(duì) node0-privateIP 和 進(jìn)行了硬編碼，但輸入了等級(jí)和 local_rank 輸入為 arg [1]和 arg [ 2]命令行參數(shù)。 上傳后，在每個(gè)節(jié)點(diǎn)中打開(kāi)兩個(gè) ssh 終端。

  • 在 node0 的第一個(gè)終端上，運(yùn)行$ python main.py 0 0
  • 在 node0 的第二個(gè)終端上運(yùn)行$ python main.py 1 1
  • 在節(jié)點(diǎn) 1 的第一個(gè)終端上，運(yùn)行$ python main.py 2 0
  • 在 node1 的第二個(gè)終端上運(yùn)行$ python main.py 3 1

在打印“ Initialize Model…”之后，程序?qū)?dòng)并等待所有四個(gè)進(jìn)程加入該進(jìn)程組。 請(qǐng)注意，第一個(gè)參數(shù)沒(méi)有重復(fù)，因?yàn)檫@是該過(guò)程的唯一全局等級(jí)。 重復(fù)第二個(gè)參數(shù)，因?yàn)檫@是在節(jié)點(diǎn)上運(yùn)行的進(jìn)程的本地等級(jí)。 如果在每個(gè)節(jié)點(diǎn)上運(yùn)行nvidia-smi，則將在每個(gè)節(jié)點(diǎn)上看到兩個(gè)進(jìn)程，一個(gè)在 GPU0 上運(yùn)行，一個(gè)在 GPU1 上運(yùn)行。

我們現(xiàn)在已經(jīng)完成了分布式訓(xùn)練示例！ 希望您能看到如何使用本教程來(lái)幫助您在自己的數(shù)據(jù)集上訓(xùn)練自己的模型，即使您沒(méi)有使用完全相同的分布式環(huán)境。 如果您使用的是 AWS，請(qǐng)不要忘記按一下來(lái)關(guān)閉您的節(jié)點(diǎn)，否則在月底可能會(huì)發(fā)現(xiàn)一張大筆的賬單。

接下來(lái)要去哪里

腳本的總運(yùn)行時(shí)間：(0 分鐘 0.000 秒）

Download Python source code: aws_distributed_training_tutorial.py Download Jupyter notebook: aws_distributed_training_tutorial.ipynb


以上內(nèi)容是否對(duì)您有幫助：
在線筆記
App下載
App下載

掃描二維碼

下載編程獅App

公眾號(hào)
微信公眾號(hào)

編程獅公眾號(hào)