最常見的深度學習框架應該是TensorFlow、Pytorch、Keras，但是這些框架在面向大規模模型的時候都不是很方便。

比如Pytorch的分布式并行計算框架（Distributed Data Parallel，簡稱DDP），它也僅僅是能將數據并行，放到各個GPU的模型上進行訓練。

也就是說，DDP的應用場景在你的模型大小大于顯卡顯存大小時，它就無法使用了，除非你自己再將模型參數拆散分散到各個GPU上。

今天要給大家介紹的DeepSpeed，它就能實現這個拆散功能，它通過將模型參數拆散分布到各個GPU上，以實現大型模型的計算，彌補了DDP的缺點，非常方便，這也就意味著我們能用更少的GPU訓練更大的模型，而且不受限于顯存。

DeepSpeed入門并不簡單，盡管是微軟開源的框架，文檔卻寫的一般，缺少條理性，也沒有從零到一的使用示例。下面我就簡單介紹一下怎么使用DeepSpeed這個框架。

1.準備

開始之前，你要確保Python和pip已經成功安裝在電腦上，如果沒有，可以訪問這篇文章：超詳細Python安裝指南 進行安裝。

**(可選1) **如果你用Python的目的是數據分析，可以直接安裝Anaconda：Python數據分析與挖掘好幫手—Anaconda，它內置了Python和pip.

**(可選2) **此外，推薦大家用VSCode編輯器，它有許多的優點：Python 編程的最好搭檔—VSCode 詳細指南。

請選擇以下任一種方式輸入命令安裝依賴 ：

1. Windows 環境 打開 Cmd (開始-運行-CMD)。
2. MacOS 環境 打開 Terminal (command+空格輸入Terminal)。
3. 如果你用的是 VSCode編輯器 或 Pycharm，可以直接使用界面下方的Terminal.

pip install deepspeed

此外，你還需要下載 Pytorch，在官網選擇自己對應的系統版本和環境，按照指示安裝即可：

https://pytorch.org/get-started/locally/

2.使用 DeepSpeed

使用DeepSpeed其實和寫一個pytorch模型只有部分區別，一開始的流程是一樣的。

2.1 載入數據集：

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()

這里我寫了一個非常簡單的模型作測試。

2.3 初始化Deepspeed

DeepSpeed 通過輸入參數來啟動訓練，因此需要使用argparse解析參數：

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

此外，模型初始化的時候除了參數，還需要model及其parameters，還有訓練集：

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 訓練邏輯

下面的部分和我們平時訓練模型是幾乎一樣的代碼，請注意 local_rank 是你不需要管的參數，在后面啟動模型訓練的時候，DeepSpeed會自動給這個參數賦值。

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 測試邏輯

模型測試和模型訓練的邏輯類似：

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 編寫模型參數

在當前目錄下新建一個 config.json 里面寫好我們的調優器、訓練batch等參數：

{
   "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
 }

完整的開發流程就結束了，可以看到其實和我們平時使用pytorch開發模型的區別不大，就是在初始化的時候使用 DeepSpeed，并以輸入參數的形式初始化。

完整代碼可以在Python實用寶典后臺回復 **Deepspeed **下載。

3. 測試代碼

現在就來測試我們上面的代碼能不能正常運行。

在這里，我們需要用環境變量控制使用的GPU，比如我的機器有10張GPU，我只使用6, 7, 8, 9號GPU，輸入命令：

export CUDA_VISIBLE_DEVICES="6,7,8,9"

然后開始運行代碼：

deepspeed test.py --deepspeed_config config.json

看到下面的輸出說明開始正常運行，在下載數據了：

開始訓練的時候 DeepSpeed 通常會打印更多的訓練細節供用戶監控，包括訓練設置、性能統計和損失趨勢，效果類似于：

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 %

當你運行到最后，出現了這樣的輸出，恭喜你，完成了你的第一個 DeepSpeed 模型，可以開始你的大規模訓練之路了。