亚洲国产精品狼友中文久久久,亚洲电影一区二区,2020国自产拍精品高潮

原文： PyTorch 單機(jī)模型并行最佳實(shí)踐

注意

作者：申力

模型并行在分布式訓(xùn)練技術(shù)中被廣泛使用。先前的帖子已經(jīng)解釋了如何使用 DataParallel 在多個(gè) GPU 上訓(xùn)練神經(jīng)網(wǎng)絡(luò)；此功能將相同的模型復(fù)制到所有 GPU，其中每個(gè) GPU 消耗輸入數(shù)據(jù)的不同分區(qū)。盡管它可以極大地加快訓(xùn)練過(guò)程，但不適用于某些模型太大而無(wú)法放入單個(gè) GPU 的用例。這篇文章展示了如何通過(guò)使用模型并行解決該問(wèn)題，與DataParallel相比，該模型將單個(gè)模型拆分到不同的 GPU 上，而不是在每個(gè) GPU 上復(fù)制整個(gè)模型(具體來(lái)說(shuō)，假設(shè)模型m包含 10 層：使用DataParallel時(shí)，每個(gè) GPU 都具有這 10 層中每個(gè)層的副本，而當(dāng)在兩個(gè) GPU 上并行使用模型時(shí)，每個(gè) GPU 可以承載 5 層）。

模型并行化的高級(jí)思想是將模型的不同子網(wǎng)放置在不同的設(shè)備上，并相應(yīng)地實(shí)現(xiàn)forward方法以在設(shè)備之間移動(dòng)中間輸出。由于模型的一部分僅在任何單個(gè)設(shè)備上運(yùn)行，因此一組設(shè)備可以共同為更大的模型服務(wù)。在本文中，我們不會(huì)嘗試構(gòu)建龐大的模型并將其壓縮到有限數(shù)量的 GPU 中。相反，本文著重展示并行模型的概念。讀者可以將這些想法應(yīng)用到實(shí)際應(yīng)用中。

Note

對(duì)于模型跨越多個(gè)服務(wù)器的分布式模型并行訓(xùn)練，請(qǐng)參見(jiàn)分布式 RPC 框架入門，以獲取示例和詳細(xì)信息。

基本用法

讓我們從包含兩個(gè)線性層的玩具模型開(kāi)始。要在兩個(gè) GPU 上運(yùn)行此模型，只需將每個(gè)線性層放在不同的 GPU 上，然后移動(dòng)輸入和中間輸出以匹配層設(shè)備。

import torch
import torch.nn as nn
import torch.optim as optim
class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')
    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

請(qǐng)注意，除了五個(gè)to(device)調(diào)用將線性層和張量放置在適當(dāng)?shù)脑O(shè)備上之外，上述ToyModel看起來(lái)非常類似于在單個(gè) GPU 上實(shí)現(xiàn)它的方式。那是模型中唯一需要更改的地方。 backward()和torch.optim將自動(dòng)處理漸變，就像模型在一個(gè) GPU 上一樣。調(diào)用損失函數(shù)時(shí)，只需確保標(biāo)簽與輸出位于同一設(shè)備上。

model = ToyModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

將模型并行應(yīng)用于現(xiàn)有模塊

只需進(jìn)行幾行更改，就可以在多個(gè) GPU 上運(yùn)行現(xiàn)有的單 GPU 模塊。以下代碼顯示了如何將torchvision.models.reset50()分解為兩個(gè) GPU。這個(gè)想法是繼承現(xiàn)有的ResNet模塊，并在構(gòu)建過(guò)程中將層拆分為兩個(gè) GPU。然后，通過(guò)相應(yīng)地移動(dòng)中間輸出，覆蓋forward方法來(lái)縫合兩個(gè)子網(wǎng)。

from torchvision.models.resnet import ResNet, Bottleneck
num_classes = 1000
class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)
        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,
            self.layer1,
            self.layer2
        ).to('cuda:0')
        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')
        self.fc.to('cuda:1')
    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

對(duì)于模型太大而無(wú)法放入單個(gè) GPU 的情況，上述實(shí)現(xiàn)解決了該問(wèn)題。但是，您可能已經(jīng)注意到，如果您的模型合適，它將比在單個(gè) GPU 上運(yùn)行它要慢。這是因?yàn)樵谌魏螘r(shí)間點(diǎn)，兩個(gè) GPU 中只有一個(gè)在工作，而另一個(gè)在那兒什么也沒(méi)做。由于中間輸出需要在layer2和layer3之間從cuda:0復(fù)制到cuda:1，因此性能進(jìn)一步惡化。

讓我們進(jìn)行實(shí)驗(yàn)以更定量地了解執(zhí)行時(shí)間。在此實(shí)驗(yàn)中，我們通過(guò)運(yùn)行隨機(jī)輸入和標(biāo)簽來(lái)訓(xùn)練ModelParallelResNet50和現(xiàn)有的torchvision.models.reset50()。訓(xùn)練后，模型將不會(huì)產(chǎn)生任何有用的預(yù)測(cè)，但是我們可以對(duì)執(zhí)行時(shí)間有一個(gè)合理的了解。

import torchvision.models as models
num_batches = 3
batch_size = 120
image_w = 128
image_h = 128
def train(model):
    model.train(True)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)
    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)
    for _ in range(num_batches):
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)
        # run forward pass
        optimizer.zero_grad()
        outputs = model(inputs.to('cuda:0'))
        # run backward pass
        labels = labels.to(outputs.device)
        loss_fn(outputs, labels).backward()
        optimizer.step()

上面的train(model)方法使用nn.MSELoss作為損失函數(shù)，并使用optim.SGD作為優(yōu)化器。它模擬了對(duì)128 X 128圖像的訓(xùn)練，這些圖像分為 3 批，每批包含 120 張圖像。然后，我們使用timeit來(lái)運(yùn)行train(model)方法 10 次，并繪制帶有標(biāo)準(zhǔn)偏差的執(zhí)行時(shí)間。

import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit
num_repeat = 10
stmt = "train(model)"
setup = "model = ModelParallelResNet50()"
## globals arg is only available in Python 3\. In Python 2, use the following
## import __builtin__
## __builtin__.__dict__.update(locals())
mp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)
setup = "import torchvision.models as models;" + \
        "model = models.resnet50(num_classes=num_classes).to('cuda:0')"
rn_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)
def plot(means, stds, labels, fig_name):
    fig, ax = plt.subplots()
    ax.bar(np.arange(len(means)), means, yerr=stds,
           align='center', alpha=0.5, ecolor='red', capsize=10, width=0.6)
    ax.set_ylabel('ResNet50 Execution Time (Second)')
    ax.set_xticks(np.arange(len(means)))
    ax.set_xticklabels(labels)
    ax.yaxis.grid(True)
    plt.tight_layout()
    plt.savefig(fig_name)
    plt.close(fig)
plot([mp_mean, rn_mean],
     [mp_std, rn_std],
     ['Model Parallel', 'Single GPU'],
     'mp_vs_rn.png')

結(jié)果表明，模型并行實(shí)現(xiàn)的執(zhí)行時(shí)間比現(xiàn)有的單 GPU 實(shí)現(xiàn)長(zhǎng)4.02/3.75-1=7%。因此，我們可以得出結(jié)論，在 GPU 之間來(lái)回復(fù)制張量大約有 7％的開(kāi)銷。有改進(jìn)的余地，因?yàn)槲覀冎纼蓚€(gè) GPU 之一在整個(gè)執(zhí)行過(guò)程中處于空閑狀態(tài)。一種選擇是將每個(gè)批次進(jìn)一步劃分為拆分管道，這樣，當(dāng)一個(gè)拆分到達(dá)第二個(gè)子網(wǎng)時(shí)，可以將下一個(gè)拆分饋入第一個(gè)子網(wǎng)。這樣，兩個(gè)連續(xù)的拆分可以在兩個(gè) GPU 上同時(shí)運(yùn)行。

通過(guò)流水線輸入加速

在以下實(shí)驗(yàn)中，我們將每批次120張圖像，進(jìn)一步劃分為 20 張圖像的均分。當(dāng) PyTorch 異步啟動(dòng) CUDA 操作時(shí)，該實(shí)現(xiàn)無(wú)需生成多個(gè)線程即可實(shí)現(xiàn)并發(fā)。

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size
    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []
        for s_next in splits:
            # A. s_prev runs on cuda:1
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
            # B. s_next runs on cuda:0, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')
        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
        return torch.cat(ret)
setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)
plot([mp_mean, rn_mean, pp_mean],
     [mp_std, rn_std, pp_std],
     ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
     'mp_vs_rn_vs_pp.png')

請(qǐng)注意，設(shè)備到設(shè)備的張量復(fù)制操作在源設(shè)備和目標(biāo)設(shè)備上的當(dāng)前流上同步。如果創(chuàng)建多個(gè)流，則必須確保復(fù)制操作正確同步。在完成復(fù)制操作之前寫入源張量或讀取/寫入目標(biāo)張量可能導(dǎo)致不確定的行為。上面的實(shí)現(xiàn)僅在源設(shè)備和目標(biāo)設(shè)備上都使用默認(rèn)流，因此沒(méi)有必要強(qiáng)制執(zhí)行其他同步。

實(shí)驗(yàn)結(jié)果表明，對(duì)并行 ResNet50 進(jìn)行建模的流水線輸入可將訓(xùn)練過(guò)程大約加快3.75/2.51-1=49%的速度。距理想的 100％加速仍然相去甚遠(yuǎn)。由于我們?cè)诠艿啦⑿袑?shí)現(xiàn)中引入了新參數(shù)split_sizes，因此尚不清楚新參數(shù)如何影響整體訓(xùn)練時(shí)間。直觀地講，使用較小的split_size會(huì)導(dǎo)致許多小的 CUDA 內(nèi)核啟動(dòng)，而使用較大的split_size會(huì)導(dǎo)致在第一次和最后一次拆分期間出現(xiàn)較長(zhǎng)的空閑時(shí)間。都不是最佳選擇。對(duì)于此特定實(shí)驗(yàn)，可能會(huì)有最佳的split_size配置。讓我們嘗試通過(guò)使用幾個(gè)不同的split_size值進(jìn)行實(shí)驗(yàn)來(lái)找到它。

means = []
stds = []
split_sizes = [1, 3, 5, 8, 10, 12, 20, 40, 60]
for split_size in split_sizes:
    setup = "model = PipelineParallelResNet50(split_size=%d)" % split_size
    pp_run_times = timeit.repeat(
        stmt, setup, number=1, repeat=num_repeat, globals=globals())
    means.append(np.mean(pp_run_times))
    stds.append(np.std(pp_run_times))
fig, ax = plt.subplots()
ax.plot(split_sizes, means)
ax.errorbar(split_sizes, means, yerr=stds, ecolor='red', fmt='ro')
ax.set_ylabel('ResNet50 Execution Time (Second)')
ax.set_xlabel('Pipeline Split Size')
ax.set_xticks(split_sizes)
ax.yaxis.grid(True)
plt.tight_layout()
plt.savefig("split_size_tradeoff.png")
plt.close(fig)

結(jié)果表明，將split_size設(shè)置為 12 可獲得最快的訓(xùn)練速度，從而導(dǎo)致3.75/2.43-1=54%加速。仍有機(jī)會(huì)進(jìn)一步加快訓(xùn)練過(guò)程。例如，對(duì)cuda:0的所有操作都放在其默認(rèn)流上。這意味著下一個(gè)拆分的計(jì)算不能與上一個(gè)拆分的復(fù)制操作重疊。但是，由于上一個(gè)和下一個(gè)拆分是不同的張量，因此將一個(gè)計(jì)算與另一個(gè)副本重疊是沒(méi)有問(wèn)題的。實(shí)現(xiàn)需要在兩個(gè) GPU 上使用多個(gè)流，并且不同的子網(wǎng)結(jié)構(gòu)需要不同的流管理策略。由于沒(méi)有通用的多流解決方案適用于所有模型并行用例，因此在本教程中將不再討論。

注意：

這篇文章顯示了幾個(gè)性能指標(biāo)。當(dāng)您在自己的計(jì)算機(jī)上運(yùn)行相同的代碼時(shí)，您可能會(huì)看到不同的數(shù)字，因?yàn)榻Y(jié)果取決于底層的硬件和軟件。為了使您的環(huán)境獲得最佳性能，一種正確的方法是首先生成曲線以找出最佳分割尺寸，然后將該分割尺寸用于管道輸入。

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

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

PyTorch 單機(jī)模型并行最佳實(shí)踐

基本用法

將模型并行應(yīng)用于現(xiàn)有模塊

通過(guò)流水線輸入加速

推薦文章

推薦教程

推薦課程