3 совета для ускорения нейронных сетей

Зачастую, сталкиваясь с задачей ускорения нейросети, инженер теряется в спектре методов, архитектур и советов на Medium. Мы выделили 3 совета, которые упростят ваше погружение в тему.

Совет №1: выбирайте модель с учетом баланса между скоростью и качеством

Представьте, что вы используете ResNet-50, обученный на датасете ImageNet. Находите метод прунинга, который показывает, что при ускорении на 70% по FLOPs точность сети будет 71.18 (самая последняя строка).

Да это ж круто! Но, давайте посмотрим, что получим в абсолютных единицах.

Количество операций сложения-умножения, или FLOPs (Floating Point Operations), или MAC (Multiply-Add Cumulation) в ResNet-50 равняется 4144.85 * 10^6. Количество операций сложения-умножения в MobileNet-v2 — 300 * 10^6.

При ускорении мы получили (1 — 0.3) * 4144.85 * 10^6 = 2901.395 * 10^6 операций.

Точность ускоренной сети — 71.18, точность MobileNet-v2 — 72.

В итоге, получили модель c меньшей точностью и большей вычислительной сложностью. А значит потратили вычислительные ресурсы и время впустую!

Совет №2: профилируйте пайплайн, прежде чем ускорять

Представьте ситуацию: сеть поджали, а время работы приложения, частью которого она является, не изменилось. При этом архитектура изменилась. В чем может быть причина?

Как вы могли догадаться, в данном случае инференс сети не является самой длительной частью пайплайна. Но как узнать это заранее? Тут на помощь приходит PyTorch Profiler.

Например, есть следующая нейронная сеть из одного слоя:

class MyModule(nn.Module):
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super(MyModule, self) .__init__()
        self.linear = nn.Linear(in_features, out_features, bias)

    def forward (self, input, mask):
        with profiler.record_function("LINEAR PASS"):
            out = self.linear (input)

        with profiler.record function ("MASK INDICES"):
            threshold = out. sum(axis=1).mean().item)
            hi_idk = np.argwhere(mask.cpu().numpy() > threshold)
            hi_idx = torch.from_numpy(hi_idk).cuda()

        return out, hi_idx

Создаем её инстанс и «прогреваем» GPU:

model = MyModule(500, 10).cuda()
input = torch.rand(128, 500).cuda()
mask = torch.rand((500, 500, 500), dtype=torch.double).cuda()
# warm-up
model (input, mask)

Обратите внимание: всегда следует «прогревать» (warm-up) свою карточку перед вычислениями.

Инференс сети представлен ниже:

with profiler.profile(with_stack=True, profile_memory=True) as prof:
    out, idx = model (input, mask)

Если вывести результаты профайлера, то можно увидеть, что большую часть времени занимает копирование памяти с CPU на GPU (cudaMemcpuAsync).

Поэтому для ускорения пайплайна этот перенос необходимо сократить.

Совет №3: используйте фьюзинг слоев

В нейронной сети есть операции, которые можно объединить в одну. Такой метод иногда называют фьюзингом (fusing — слияние) или сплавлением. Рассмотрим самый частый пример со сверткой (Conv) и Batch Normalization (BN).

Напомним, что BN подстраивает распределения, которые были внутри сети, так, чтобы те имели среднее \( \beta \) и дисперсию \( \gamma \). Математически, это выглядит так:

\( \mu_{B} = \frac{1}{m} \sum_{i=1}^{m} \textbf{x}^{(i)} \)

\( \sigma^2_{B} = \frac{1}{m} \sum_{i=1}^{m} (\textbf{x}^{(i)} – \mu)^2 \)

\( \hat{\textbf{x}}^{(i)} = \frac{\textbf{x}^{(i)} – \mu_{B}}{\sqrt{\sigma^2_{B} + \epsilon}} \)

\( \textbf{y}^{(i)} = \gamma \hat{\textbf{x}}^{(i)} + \beta \)

Важно помнить, что \( \mu_{B}, \sigma^{2}_{B} \) на инференсе не вычисляются, а используются уже посчитанные в процессе обучения. Таким образом, можно просто раскрыть скобки в последней формуле:

\( \textbf{y}^{(i)} = \gamma \hat{\textbf{x}}^{(i)} + \beta = \frac{\hat{\textbf{x}}^{(i)} – \mu_{B}}{\sqrt{\sigma^2_{B} + \epsilon}}*\gamma + \beta = \frac{\hat{\textbf{x}}^{(i)}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} + \beta — \frac{\mu_{B}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} \)

Далее вспоминаем, что \( x^{(i)} \) — не просто какой-то тензор, а результат свертки или применения полносвязного слоя. Предлагается просто “сплавить” веса этого слоя и масштабирующий коэффициент, который получился в предыдущей формуле:

\( x^{(i)} = z^{(i)}*W \)

\( \textbf{y}^{(i)} = \frac{\textbf{x}^{(i)}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} + \beta — \frac{\mu_{B}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} = z^{(i)}\frac{W\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} + \beta — \frac{\mu_{B}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} = z^{(i)}\hat{W} + \hat{b} \)

\( \hat{W} = \frac{W\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} \)

\( \hat{b}= \beta — \frac{\mu_{B}\gamma}{\sqrt{\sigma^2_{B} + \epsilon}} \)

В итоге, получим абсолютно идентичный результат при уменьшении количества вычислений!

Применить сплавление из коробки можно на этапе сохранения в onnx — основной формат для сохранения моделей перед инференсом на любом из фреймворков (TensorRT, OpenVINO и прочие). Нужные аргументы в функции выделены жирным шрифтом.

# Input to the model
x = torch. rand(batch_size, 1, 224, 224, requires_grad=True)
torch_out = torch_model (x)

# Export the model
torch.onnx. export(
    torch_model,            # model being run
    x,                      # model input (or a tuple for multiple inputs)
    "super_resolution.onnx",# where to save the model (can be a file or file-like object)
    export_params=True,            # store the trained parameter weights inside the model file
    opset version=10,              # the ONNX version to export the model to
    **do_constant_folding=True,      # whether to execute constant folding for optimization**
    input_names = ['input'],       # the model's input names
    output_names = ['output'],     # the model's output names
)

Если же вы делаете инференс на PyTorch (что, кстати, зря), то можете воспользоваться скриптом, который подготовили держатели репозитория.

Заключение

Мы дали 3 рекомендации, которые могут быть полезны при ускорении нейронных сетей. На первый взгляд, советы могут показаться очевидными, но их часто упускают.

При этом не соблюдая первые два пункта можно потратить огромное количество времени впустую, а значит, затянуть сроки и, возможно, провалить проект.