大規(guī)模深度學(xué)習(xí)框架 DeepSpeed 使用指南
最常見的深度學(xué)習(xí)框架應(yīng)該是TensorFlow、Pytorch、Keras，但是這些框架在面向大規(guī)模模型的時(shí)候都不是很方便。
比如Pytorch的分布式并行計(jì)算框架（Distributed Data Parallel，簡稱DDP），它也僅僅是能將數(shù)據(jù)并行，放到各個(gè)GPU的模型上進(jìn)行訓(xùn)練。
也就是說，DDP的應(yīng)用場景在你的模型大小大于顯卡顯存大小時(shí)，它就無法使用了，除非你自己再將模型參數(shù)拆散分散到各個(gè)GPU上。
今天要給大家介紹的DeepSpeed，它就能實(shí)現(xiàn)這個(gè)拆散功能，它通過將模型參數(shù)拆散分布到各個(gè)GPU上，以實(shí)現(xiàn)大型模型的計(jì)算，彌補(bǔ)了DDP的缺點(diǎn)，非常方便，這也就意味著我們能用更少的GPU訓(xùn)練更大的模型，而且不受限于顯存。
DeepSpeed入門并不簡單，盡管是微軟開源的框架，文檔卻寫的一般，缺少條理性，也沒有從零到一的使用示例。下面我就簡單介紹一下怎么使用DeepSpeed這個(gè)框架。
1.準(zhǔn)備
開始之前，你要確保Python和pip已經(jīng)成功安裝在電腦上，如果沒有，可以訪問這篇文章：超詳細(xì)Python安裝指南?進(jìn)行安裝。
(可選1)?如果你用Python的目的是數(shù)據(jù)分析，可以直接安裝Anaconda：Python數(shù)據(jù)分析與挖掘好幫手—Anaconda，它內(nèi)置了Python和pip.
(可選2)?此外，推薦大家用VSCode編輯器，它有許多的優(yōu)點(diǎn)：Python 編程的最好搭檔—VSCode 詳細(xì)指南。
請選擇以下任一種方式輸入命令安裝依賴：
1. Windows 環(huán)境 打開 Cmd (開始-運(yùn)行-CMD)。
2. MacOS 環(huán)境 打開 Terminal (command+空格輸入Terminal)。
3. 如果你用的是 VSCode編輯器 或 Pycharm，可以直接使用界面下方的Terminal.
pip install deepspeed
此外，你還需要下載 Pytorch，在官網(wǎng)選擇自己對應(yīng)的系統(tǒng)版本和環(huán)境，按照指示安裝即可：
https://pytorch.org/get-started/locally/
2.使用 DeepSpeed
使用DeepSpeed其實(shí)和寫一個(gè)pytorch模型只有部分區(qū)別，一開始的流程是一樣的。
2.1 載入數(shù)據(jù)集：
import?torch
import?torchvision
import?torchvision.transforms as?transforms
trainset = torchvision.datasets.CIFAR10(root='./data',
????????????????????????????????????????train=True,
????????????????????????????????????????download=True,
????????????????????????????????????????transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
??????????????????????????????????????????batch_size=16,
??????????????????????????????????????????shuffle=True,
??????????????????????????????????????????num_workers=2)
testset = torchvision.datasets.CIFAR10(root='./data',
???????????????????????????????????????train=False,
???????????????????????????????????????download=True,
???????????????????????????????????????transform=transform)
testloader = torch.utils.data.DataLoader(testset,
?????????????????????????????????????????batch_size=4,
?????????????????????????????????????????shuffle=False,
?????????????????????????????????????????num_workers=2)2.2 編寫模型：
import?torch.nn as?nn
import?torch.nn.functional as?F
class?Net(nn.Module):
????def?__init__(self):
????????super(Net, self).__init__()
????????self.conv1 = nn.Conv2d(3, 6, 5)
????????self.pool = nn.MaxPool2d(2, 2)
????????self.conv2 = nn.Conv2d(6, 16, 5)
????????self.fc1 = nn.Linear(16?* 5?* 5, 120)
????????self.fc2 = nn.Linear(120, 84)
????????self.fc3 = nn.Linear(84, 10)
????def?forward(self, x):
????????x = self.pool(F.relu(self.conv1(x)))
????????x = self.pool(F.relu(self.conv2(x)))
????????x = x.view(-1, 16?* 5?* 5)
????????x = F.relu(self.fc1(x))
????????x = F.relu(self.fc2(x))
????????x = self.fc3(x)
????????return?x
net = Net()
criterion = nn.CrossEntropyLoss()這里我寫了一個(gè)非常簡單的模型作測試。
2.3 初始化Deepspeed
DeepSpeed 通過輸入?yún)?shù)來啟動(dòng)訓(xùn)練，因此需要使用argparse解析參數(shù)：
import?argparse
def?add_argument():
????parser = argparse.ArgumentParser(description='CIFAR')
????parser.add_argument('-b',
????????????????????????'--batch_size',
????????????????????????default=32,
????????????????????????type=int,
????????????????????????help='mini-batch size (default: 32)')
????parser.add_argument('-e',
????????????????????????'--epochs',
????????????????????????default=30,
????????????????????????type=int,
????????????????????????help='number of total epochs (default: 30)')
????parser.add_argument('--local_rank',
????????????????????????type=int,
????????????????????????default=-1,
????????????????????????help='local rank passed from distributed launcher')
????parser.add_argument('--log-interval',
????????????????????????type=int,
????????????????????????default=2000,
????????????????????????help="output logging information at a given interval")
????parser = deepspeed.add_config_arguments(parser)
????args = parser.parse_args()
????return?args此外，模型初始化的時(shí)候除了參數(shù)，還需要model及其parameters，還有訓(xùn)練集：
args = add_argument()
net = Net()
parameters = filter(lambda?p: p.requires_grad, net.parameters())
model_engine, optimizer, trainloader, __ = deepspeed.initialize(
????args=args, model=net, model_parameters=parameters, training_data=trainset)2.4 訓(xùn)練邏輯
下面的部分和我們平時(shí)訓(xùn)練模型是幾乎一樣的代碼，請注意 local_rank 是你不需要管的參數(shù)，在后面啟動(dòng)模型訓(xùn)練的時(shí)候，DeepSpeed會(huì)自動(dòng)給這個(gè)參數(shù)賦值。
for?epoch in?range(2):
????running_loss = 0.0
????for?i, data in?enumerate(trainloader):
????????inputs, labels = data[0].to(model_engine.local_rank), data[1].to(
????????????model_engine.local_rank)
????????outputs = model_engine(inputs)
????????loss = criterion(outputs, labels)
????????model_engine.backward(loss)
????????model_engine.step()
????????# print statistics
????????running_loss += loss.item()
????????if?i % args.log_interval == (args.log_interval - 1):
????????????print('[%d, %5d] loss: %.3f'?% (epoch + 1, i + 1, running_loss / args.log_interval))
????????????running_loss = 0.02.5 測試邏輯
模型測試和模型訓(xùn)練的邏輯類似：
correct = 0
total = 0
with?torch.no_grad():
????for?data in?testloader:
????????images, labels = data
????????outputs = net(images.to(model_engine.local_rank))
????????_, predicted = torch.max(outputs.data, 1)
????????total += labels.size(0)
????????correct += (predicted == labels.to(
????????????model_engine.local_rank)).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%'?%
??????(100?* correct / total))2.6 編寫模型參數(shù)
在當(dāng)前目錄下新建一個(gè) config.json 里面寫好我們的調(diào)優(yōu)器、訓(xùn)練batch等參數(shù)：
{
???"train_batch_size": 4,
???"steps_per_print": 2000,
???"optimizer": {
?????"type": "Adam",
?????"params": {
???????"lr": 0.001,
???????"betas": [
?????????0.8,
?????????0.999
???????],
???????"eps": 1e-8,
???????"weight_decay": 3e-7
?????}
???},
???"scheduler": {
?????"type": "WarmupLR",
?????"params": {
???????"warmup_min_lr": 0,
???????"warmup_max_lr": 0.001,
???????"warmup_num_steps": 1000
?????}
???},
???"wall_clock_breakdown": false
?}
完整的開發(fā)流程就結(jié)束了，可以看到其實(shí)和我們平時(shí)使用pytorch開發(fā)模型的區(qū)別不大，就是在初始化的時(shí)候使用 DeepSpeed，并以輸入?yún)?shù)的形式初始化。
完整代碼可以在Python實(shí)用寶典后臺(tái)回復(fù)?Deepspeed?下載。
3. 測試代碼
現(xiàn)在就來測試我們上面的代碼能不能正常運(yùn)行。
在這里，我們需要用環(huán)境變量控制使用的GPU，比如我的機(jī)器有10張GPU，我只使用6, 7, 8, 9號GPU，輸入命令：
export?CUDA_VISIBLE_DEVICES="6,7,8,9"然后開始運(yùn)行代碼：
deepspeed test.py --deepspeed_config config.json看到下面的輸出說明開始正常運(yùn)行，在下載數(shù)據(jù)了：
開始訓(xùn)練的時(shí)候 DeepSpeed 通常會(huì)打印更多的訓(xùn)練細(xì)節(jié)供用戶監(jiān)控，包括訓(xùn)練設(shè)置、性能統(tǒng)計(jì)和損失趨勢，效果類似于：
worker-0: [INFO 2020-02-06 20:35:23] 0/24550, SamplesPerSec=1284.4954513975558
worker-0: [INFO 2020-02-06 20:35:23] 0/24600, SamplesPerSec=1284.384033658866
worker-0: [INFO 2020-02-06 20:35:23] 0/24650, SamplesPerSec=1284.4433482972925
worker-0: [INFO 2020-02-06 20:35:23] 0/24700, SamplesPerSec=1284.4664449792422
worker-0: [INFO 2020-02-06 20:35:23] 0/24750, SamplesPerSec=1284.4950124403447
worker-0: [INFO 2020-02-06 20:35:23] 0/24800, SamplesPerSec=1284.4756105952233
worker-0: [INFO 2020-02-06 20:35:24] 0/24850, SamplesPerSec=1284.5251526215386
worker-0: [INFO 2020-02-06 20:35:24] 0/24900, SamplesPerSec=1284.531217073863
worker-0: [INFO 2020-02-06 20:35:24] 0/24950, SamplesPerSec=1284.5125323220368
worker-0: [INFO 2020-02-06 20:35:24] 0/25000, SamplesPerSec=1284.5698818883018
worker-0: Finished Training
worker-0: GroundTruth: cat ship ship plane
worker-0: Predicted: cat car car plane
worker-0: Accuracy of the network on the 10000 test images: 57 %當(dāng)你運(yùn)行到最后，出現(xiàn)了這樣的輸出，恭喜你，完成了你的第一個(gè) DeepSpeed 模型，可以開始你的大規(guī)模訓(xùn)練之路了。
我們的文章到此就結(jié)束啦，如果你喜歡今天的Python 實(shí)戰(zhàn)教程，請持續(xù)關(guān)注Python實(shí)用寶典。
有任何問題，可以在公眾號后臺(tái)回復(fù)：加群，回答相應(yīng)紅字驗(yàn)證信息，進(jìn)入互助群詢問。
原創(chuàng)不易，希望你能在下面點(diǎn)個(gè)贊和在看支持我繼續(xù)創(chuàng)作，謝謝！
點(diǎn)擊下方閱讀原文可獲得更好的閱讀體驗(yàn)
Python實(shí)用寶典?(pythondict.com)
不只是一個(gè)寶典
歡迎關(guān)注公眾號：Python實(shí)用寶典