Что делать, если ваша модель работает медленнее, чем ожидалось

Шаги, описанные в этой статье, также задокументированы в этой проблеме GitHub

ONNX Runtime — кроссплатформенный ускоритель машинного обучения для логического вывода и обучения. Он обеспечивает единый стандартизированный формат для выполнения моделей машинного обучения.

Чтобы дать представление о широте поддержки, на изображении ниже показаны все текущие платформы сборки.

ONNX подает большие надежды и является отличным проектом. Это расширяет возможности выполнения модели. Но это не всегда просто.

В этой статье будет рассмотрен случай экспорта модели PyTorch в ONNX и что было сделано для повышения производительности графического процессора.

Фон

txtai недавно добавил поддержку моделей преобразования текста в речь. Выбранными моделями были модели PyTorch в ESPnet, экспортированные в ONNX с использованием espnet_onnx. Эти модели ONNX доступны на Hugging Face Hub.

Подробнее об этом можно прочитать по ссылке ниже.



Генерация текста в речь
Эта статья является частью серии руководств по txtai, платформе семантического поиска на основе искусственного интеллекта. Модели преобразования текста в речь (TTS)…neuml.hashnode.dev



Проблема

В следующих двух разделах кода показан минимальный пример выполнения вывода с использованием ESPnet напрямую (PyTorch) и запуска той же модели через ONNX.

Сначала код, использующий ESPnet напрямую и PyTorch.

import time

from espnet2.bin.tts_inference import Text2Speech

model = "espnet/kan-bayashi_ljspeech_vits"
model = Text2Speech.from_pretrained(model, device="cuda")

def run(text):
    start = time.time()
    output = model(text)
    speech = output["wav"].cpu().numpy()
    print("Time:", time.time() - start)

Запуск приведенного выше кода после прогрева (второй запуск) приводит к следующему:

>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 0.16863441467285156

Затем модель запускается с ONNX.

import time

import onnxruntime
import yaml

from ttstokenizer import TTSTokenizer

with open("ljspeech-vits-onnx/config.yaml", "r", encoding="utf-8") as f:
    config = yaml.safe_load(f)

tokenizer = TTSTokenizer(config["token"]["list"])

model = onnxruntime.InferenceSession(
    "ljspeech-vits-onnx/model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

def run(text):
    print(text)

    # Tokenize text to phoneme token ids
    inputs = tokenizer(text)

    start = time.time()
    outputs = model.run(None, {"text": inputs})
    speech = outputs[0]
    print("Time:", time.time() - start)

Запуск этого кода выше после прогрева (второй запуск) приводит к следующему:

>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 1.134232521057129

0,17 с в PyTorch против 1,13 с в ONNX, оба с включенными графическими процессорами. PyTorch почти в 7 раз быстрее.

Посмотрим, сможем ли мы добраться до сути этого.

Попытка № 1 — привязка ввода-вывода

После пары поисковых запросов в Интернете по запросу PyTorch vs ONNX slow чаще всего всплывал вопрос, связанный с передачей данных от процессора к графическому процессору. Хотя входные данные для этой модели невелики, выходные данные wav могут легко достигать 100 КБ+. Так что это может быть оно. Подробнее о передаче данных в ONNX можно прочитать в Руководстве по API.

Мы изменим наш метод запуска следующим образом и перезапустим его.

def run(text):
    print(text)

    # Tokenize text to phoneme token ids
    inputs = tokenizer(text)

    io_binding = model.io_binding()
    io_binding.bind_cpu_input("text", inputs)
    io_binding.bind_output("wav", "cuda")

    start = time.time()

    # Run model
    model.run_with_iobinding(io_binding)
    outputs = io_binding.copy_outputs_to_cpu()
    print("Time:", time.time() - start)
>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 1.1370353698730469

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

Включить профилирование

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

opts = onnxruntime.SessionOptions()
opts.enable_profiling = True

model = onnxruntime.InferenceSession(
    "ljspeech-vits-onnx/model.onnx",
     opts,
     providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
)

Файл JSON в формате onnxruntime_profile_{DATE}.json должен находиться в каталоге среды выполнения. Этот файл имеет длинный список элементов JSON с продолжительностью выполнения каждого узла.

Давайте запустим следующее, чтобы извлечь и найти самые медленные узлы.

$ jq . onnxruntime_profile.json | grep dur | cut -d ":" -f2 | sort --numeric

Мы возьмем самые большие времена и поищем их в файле onnxruntime_profile.

Первые несколько раз были связаны с полным запуском модели и инициализацией сеанса. Но после этого выделился этот.

    "cat": "Node",
    "pid": 1027615,
    "tid": 1027615,
    "dur": 313964,
    "ts": 4425685,
    "ph": "X",
    "name": "/w_1/Conv_kernel_time",

И таких было больше. Это сверточный слой.

Попытка №2 — Настройки CUDA

Давайте заглянем в документацию и посмотрим, какие опции CUDA доступны.

cudnn_conv_algo_search — вариант, который выделялся больше всего. Значение по умолчанию EXHAUSTIVE с упоминанием дорогого также показалось уместным.

Давайте попробуем изменить этот параметр и перезапустить.

opts = onnxruntime.SessionOptions()
opts.enable_profiling = True

model = onnxruntime.InferenceSession(
  "ljspeech-vits-onnx/model.onnx",
  opts,
  providers=[
    ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
    "CPUExecutionProvider"
  ]
)
>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 0.1624901294708252

Теперь среда выполнения такая же, если не немного лучше, чем PyTorch!

Вот полный код, используемый для достижения этих результатов.

import time

import onnxruntime
import yaml

from ttstokenizer import TTSTokenizer

with open("ljspeech-vits-onnx/config.yaml", "r", encoding="utf-8") as f:
    config = yaml.safe_load(f)

tokenizer = TTSTokenizer(config["token"]["list"])

model = onnxruntime.InferenceSession(
  "ljspeech-vits-onnx/model.onnx",
  providers=[
    ("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),
    "CPUExecutionProvider"
  ]
)

def run(text):
    print(text)

    # Tokenize text to phoneme token ids
    inputs = tokenizer(text)

    start = time.time()
    outputs = model.run(None, {"text": inputs})
    speech = outputs[0]
    print("Time:", time.time() - start)

Осмысление результатов

Давайте взглянем на PyTorch и посмотрим, есть ли намеки на то, что может объяснить разницу с настройками по умолчанию.

Файл Conv_v7.cpp в PyTorch имеет следующую логику, которая представляется актуальной.

  // Code reduced to these statements for clarity
  static constexpr auto DEFAULT_ALGO =
     CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_PRECOMP_GEMM;

  if (!benchmark) {
    AT_CUDNN_CHECK_WITH_SHAPES(cudnnGetConvolutionForwardAlgorithm_v7(
    )
  } else {
    AT_CUDNN_CHECK_WITH_SHAPES(cudnnFindConvolutionForwardAlgorithmEx(
    )
  }

Открытие нового Python REPL и запуск следующих шоу.

>>> import torch
>>> torch.backends.cudnn.deterministic
False
>>> torch.backends.cudnn.benchmark
False

Это означает, что cudnnGetConvolutionForwardAlgorithm_v7 является функцией по умолчанию. Даже если для параметра benchmark задано значение True, PyTorch будет кэшировать результат, если входная форма является статической. В случае этой модели входы являются динамическими. Дополнительную информацию см. в этой ветке обсуждения.

Давайте проверим это на исходном примере ESPnet PyTorch.

>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 0.17191290855407715
>>> 
>>> import torch
>>> torch.backends.cudnn.benchmark
False
>>> torch.backends.cudnn.benchmark = True
>>> torch.backends.cudnn.benchmark
True
>>> run("warmup")
>>> run("Text to speech models have recently made great strides in quality")
Time: 1.0886101722717285

Установка для torch.backends.cudnn.benchmark значения True соответствует исходному результату ONNX.

Оглядываясь назад на документацию ONNX выше, в ней говорится:

ЭВРИСТИКА (1)

облегченный эвристический поиск с использованием cudnnGetConvolutionForwardAlgorithm_v7

Похоже, что если не указано иное, PyTorch не выполняет исчерпывающий поиск. По умолчанию PyTorch будет запускать то, что ONNX называет HEURISTIC. При тестировании HEURISTIC и DEFAULT была получена одинаковая производительность.

Подведение итогов

В этой статье рассказывается, как отлаживать производительность модели ONNX, особенно при работе на графическом процессоре. Чаще всего обращают внимание на передачу данных от ЦП к графическому процессору и делают ее более эффективной с помощью привязок ввода-вывода. Но в этом случае изменение настройки CUDA было решением.

Возможно, это единичный случай с данной моделью или железом, на котором проводились тесты. Но это по крайней мере то, что можно проверить, если эта проблема возникнет в будущем!