深度學(xué)習(xí)框架DeepSpeed使用指南

最常見的深度學(xué)習(xí)框架應(yīng)該是TensorFlow、Pytorch、Keras，但是這些框架在面向大規(guī)模模型的時(shí)候都不是很方便。

比如Pytorch的分布式并行計(jì)算框架（Distributed Data Parallel，簡(jiǎn)稱DDP），它也僅僅是能將數(shù)據(jù)并行，放到各個(gè)GPU的模型上進(jìn)行訓(xùn)練。

也就是說，DDP的應(yīng)用場(chǎ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入門并不簡(jiǎn)單，盡管是微軟開源的框架，文檔卻寫的一般，缺少條理性，也沒有從零到一的使用示例。下面我就簡(jiǎn)單介紹一下怎么使用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ì)指南。

請(qǐng)選擇以下任一種方式輸入命令安裝依賴 ：

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)選擇自己對(duì)應(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è)非常簡(jiǎn)單的模型作測(cè)試。

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)練模型是幾乎一樣的代碼，請(qǐng)注意 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.0

2.5 測(cè)試邏輯

模型測(cè)試和模型訓(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. 測(cè)試代碼

現(xiàn)在就來測(cè)試我們上面的代碼能不能正常運(yùn)行。

在這里，我們需要用環(huán)境變量控制使用的GPU，比如我的機(jī)器有10張GPU，我只使用6, 7, 8, 9號(hào)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ì)和損失趨勢(shì)，效果類似于：

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)練之路了。