Автор: Александр Гончаренко, СТО enot.ai

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

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

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

Рисунок 1. Результаты работы алгоритма HAP на сети ResNet-50. Максимальное ускорение сети получается на последней строке.

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

Количество операций сложения-умножения, или 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 меньшей точностью и большей вычислительной сложностью. А значит потратили вычислительные ресурсы и время впустую!

<aside> 💡 Чтобы не терять время зря, постарайтесь заранее оценить желаемую производительность и выпишите несколько сценариев ускорения модели. Например, возможно, не стоит ускорять ResNet в 20 раз, а лучше ускорить MobileNet в 2 раза.

</aside>

Совет №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) свою карточку перед вычислениями.