Оптимизация производства с помощью PyTorch/TF, ONNX, TensorRT и LiteRT

Оптимизация производства с помощью PyTorch/TF, ONNX, TensorRT и LiteRT


Введение

Фреймворки для машинного обучения, инструменты для работы с моделями и стратегии развертывания имеют различные варианты использования в конвейере машинного обучения (ML). Они могут проявлять сильные и слабые стороны на каждом этапе разработки, обучения, оптимизации, развертывания и инференса модели.

В этой статье мы подробно сравниваем пять таких инструментов и связанных технологий: PyTorch, TensorFlow, LiteRT (ранее известный как TensorFlow Lite), TensorRT и ONNX. Мы рассматриваем их характеристики, сильные и слабые стороны, а также их роль в конвейере машинного обучения. Мы также обсудим, как различные инструменты работают вместе и поделимся общими способами их развертывания, включая примеры кода для лучшего объяснения концепций.

Основные выводы

  • Соответствие инструментов этапам: Начните с PyTorch/TF для обучения, затем используйте ONNX для переноса моделей, TensorRT для оптимизации на NVIDIA-GPU и развертывайте на мобильных устройствах/на периферии с помощью LiteRT.
  • Оптимизация по специализации: TensorRT снижает задержку и увеличивает пропускную способность на GPU; LiteRT уменьшает размер бинарных файлов и потребление памяти на устройствах.
  • План производства сантехники: TensorFlow предлагает интегрированные конвейеры (TFX/TensorFlow Serving), в то время как PyTorch обычно объединяет TorchServe и ONNX Runtime/TensorRT с большим количеством связующего кода.
  • Опирайтесь на совместимость: используйте путь PyTorch → ONNX → TensorRT для обслуживания на GPU, или TensorFlow → LiteRT для приложений на устройстве.
  • Упрощено с управляемой платформой: Платформа Linux-Console.net Gradient AI может объединять обучение и развертывание для PyTorch/TF с ONNX, TensorRT и LiteRT в одном рабочем процессе.

PyTorch: гибкая платформа для обучения и дизайн, ориентированный на исследования

PyTorch — это открытая платформа для глубокого обучения, основанная на динамическом графе вычислений (define-by-run) и Python-подобном интерфейсе программирования. PyTorch сосредоточен на идее гибкости (модели можно писать и отлаживать естественным образом, как обычный Python-код).

Хотя PyTorch по своей сути является динамичным, его высоко оптимизированная C++-бэкенд и тензорные библиотеки (такие как cuDNN для GPU) позволяют ему достигать сопоставимой или даже лучшей производительности по сравнению с фреймворком со статическим графом. Со временем PyTorch эволюционировал, чтобы поддерживать исследования, разработку и развертывание в продакшене, с такими функциями, как TorchScript (способ сериализации моделей для оптимизированного вывода в C++ или мобильных средах) и TorchServe для обслуживания моделей.

PyTorch: сильные и слабые стороны

Ниже приведен обзор того, в чем PyTorch превосходит, а где существуют компромиссы.

Сильные стороны PyTorch

  • Пайтоновский API — как обычный Python-код, его легко писать и отлаживать.
  • Динамический граф вычислений — ускоряет эксперименты и делает отладку более естественной.
  • Богатая экосистема и сообщество — видение, обработка естественного языка, инструменты и библиотеки.
  • Нативные инструменты Python — хорошо интегрируются с нативными отладчиками/профайлерами Python.
  • Улучшение производительности — JIT и AOT-компиляция PyTorch 2.x для более быстрого вывода и развертывания.
  • Сначала GPU — легкое использование ускорения CUDA; мощное распределенное обучение и развертывание на нескольких GPU.
  • Развертывания созревают — TorchScript, TorchServe и Torch-TensorRT уменьшают разрыв между исследованиями и производством.

Слабые стороны PyTorch

  • Не готово к производству «под ключ» (по сравнению с TF) — PyTorch обычно подразумевает стек TorchServe + ONNX Runtime/TensorRT + собственный CI/CD, с большим количеством вспомогательного кода и накладных операций.
  • Необходима жёсткая подготовка к поставке — Динамическая природа графов в PyTorch требует дополнительного скриптинга (TorchScript) или методов компиляции заранее (Ahead-of-Time), чтобы создавать детерминированные, воспроизводимые артефакты, для которых тесты на паритет необходимы, чтобы избежать расхождения поведения между разработкой и продакшеном.
  • Инструменты MLOps — PyTorch требует более разнородного набора инструментов. Это увеличивает объём работы по поддержке (смещение версий, обновления, исправления безопасности).
  • Большее использование ресурсов во время выполнения — Полная среда выполнения PyTorch (включая зависимости CUDA/ускорителей) достаточно тяжела для мобильных и edge-приложений.

PyTorch на практике: процесс обучения и развертывания

PyTorch в первую очередь используется на этапе разработки и обучения модели. На стадии развертывания вы можете использовать PyTorch для выполнения вывода. Например, запуск на сервере (при желании обернув модель веб-сервисом, таким как Flask, или используя TorchServe), либо конвертируя модель в более легкий формат для развертывания.

Пример PyTorch – Использование для обучения модели

В этом коде мы создаем простую полностью связанную нейронную сеть и обучаем её с использованием nn.Module и API оптимизации PyTorch. Благодаря динамическому графу PyTorch мы можем описывать тренировочный цикл на Python, где итерация — это просто выполнение прямого прохода нашей модели. После обучения мы можем сохранить веса модели. В дальнейшем мы можем либо заново загрузить модель в PyTorch, либо экспортировать её в ONNX для использования с другими средами выполнения (мы рассмотрим это позже в этом руководстве).

import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms from torch.utils.data import DataLoader  # Define a simple neural network class SimpleNet(nn.Module):     def __init__(self):         super(SimpleNet, self).__init__()         self.fc1 = nn.Linear(784, 128)         self.relu = nn.ReLU()         self.fc2 = nn.Linear(128, 10)     def forward(self, x):         x = x.view(-1, 784)  # Flatten images         x = self.relu(self.fc1(x))         return self.fc2(x)  # Prepare the training dataset and DataLoader transform = transforms.Compose([transforms.ToTensor()]) train_dataset = datasets.MNIST(root="./data", train=True, transform=transform, download=True) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)  # Model setup device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = SimpleNet().to(device) optimizer = optim.Adam(model.parameters(), lr=1e-3) criterion = nn.CrossEntropyLoss()  # Training loop for epoch in range(5):     for batch_x, batch_y in train_loader:         batch_x, batch_y = batch_x.to(device), batch_y.to(device)         optimizer.zero_grad()         outputs = model(batch_x)         loss = criterion(outputs, batch_y)         loss.backward()         optimizer.step()  print("Training completed.") torch.save(model.state_dict(), "model_weights.pth")

TensorFlow: Фреймворк и экосистема производственного уровня

TensorFlow — это еще один популярный фреймворк для глубокого обучения, изначально разработанный в Google. В ранних версиях он вводил концепцию статических графов вычислений, когда граф вычислений модели сначала определяется, а затем выполняется. Это облегчало оптимизацию всей модели и эффективное развертывание на различных платформах, но несколько ограничивало гибкость и удобство использования. В ответ на отзывы сообщества (а также под давлением конкуренции с PyTorch) TensorFlow 2.x по умолчанию принял исполнение по требованию (более динамичное, похожее на PyTorch), при этом по-прежнему предоставляя пользователям возможность использовать оптимизации статических графов с помощью декоратора tf.function и компилятора XLA.

TensorFlow также поддерживает Keras в качестве своего высокоуровневого API для определения моделей. Он предоставляет богатую экосистему инструментов для производства. Вы можете использовать его для развертывания моделей на серверах (TensorFlow Serving), на мобильных/встроенных устройствах (TensorFlow Lite, теперь LiteRT), в JavaScript (TensorFlow.js) или на специализированном оборудовании (TPU от Google).

TensorFlow: сильные и слабые стороны

TensorFlow широко используется в промышленности для обучения на больших масштабах и развёртывания в производстве. Его экосистема предоставляет решение «от конца до конца», но она имеет свои собственные особенности и компромиссы для разработчиков.

Сильные стороны TensorFlow

  • Высокая масштабируемость: Встроенная поддержка распределённого обучения на нескольких GPU и даже на нескольких машинах.
  • Готовая к производству экосистема: TensorFlow Extended включает компоненты для загрузки данных, проверки данных и обслуживания моделей для корпоративных ML-пайплайнов.
  • Оптимизированные вычисления: Выполнение статического графа (при включении) и компилятор XLA могут значительно повысить производительность.
  • Интеграция с TPU: Родная поддержка обучения и инференса на Google TPU для ускорения рабочих нагрузок.
  • Бесшовное развертывание: формат SavedModel работает с TensorFlow Serving для беспроблемного использования в производстве.
  • Готово для Edge и мобильных устройств: Конвертация в LiteRT проста для развертывания на устройстве.
  • Богатые ресурсы: Большое сообщество, множество ресурсов и предварительно обученные модели для начала переноса обучения.

Слабые стороны TensorFlow

  • Исторически крутая кривая обучения: TensorFlow 1.x требовал явного управления графом/сессиями, что могло быть запутанным для новых пользователей…
  • Сложность отладки: Использование декоратора @tf.function и графового режима в TF 2.x также не обеспечивает построчного выполнения, что может усложнять отладку.
  • Серьезная нагрузка на время выполнения: Полный TensorFlow большой и непрактичный для прямого использования на устройствах — отсюда необходимость LiteRT/TFLite.
  • Сложные пользовательские операции: Написание новых операций или ядер требует более глубоких знаний внутренней структуры TensorFlow.
  • Отставание в принятии исследовательских тенденций: Некоторые передовые слои и методы сначала появляются в PyTorch, прежде чем их перенесут в TensorFlow.
  • Менее гибко для экспериментов: Поведение статического графа может казаться жестким по сравнению с динамическим графом вычислений PyTorch во время быстрого прототипирования.

TensorFlow: от разработки до развертывания

Мы можем разработать и обучить модель с помощью TensorFlow, обычно используя API Keras, экспортировать модель и затем развернуть её.

Экспортированную модель можно использовать для вывода в серверной среде (с использованием TensorFlow Serving или TensorFlow C++ API). Для устройств периферийного типа (например, мобильных или IoT) можно преобразовать модель в формат LiteRT (TFLite) и запускать вывод непосредственно на устройстве. TensorFlow обеспечивает интеграцию с TensorRT для ускорения вывода на GPU (TF-TRT). В одном случае было зафиксировано увеличение пропускной способности вывода в 2,4 раза для модели ResNet-50 с использованием GPU NVIDIA T4 по сравнению с выполнением на стандартном GPU TensorFlow.

Пример TensorFlow – Обучение модели Keras

В приведённом ниже примере мы создаём простую полносвязную сеть для классификации изображений с использованием Keras. Мы компилируем и обучаем модель на x_train, y_train (подготовленных данных MNIST в этом примере) и сохраняем её. SavedModel („model_saved.keras“) можно загрузить позже для инференса или преобразовать для развертывания. Высокоуровневый API TensorFlow позволяет нам не создавать низкоуровневый граф вручную. Однако в фоновом режиме он может оптимизировать граф вычислений для повышения производительности при работе в продакшене.

import tensorflow as tf  # Load MNIST dataset (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()  # Normalize the pixel values to [0, 1] x_train = x_train / 255.0 x_test = x_test / 255.0  # Define a simple model using Keras (e.g., for MNIST classification) model = tf.keras.Sequential([     tf.keras.layers.Flatten(input_shape=(28, 28)),     tf.keras.layers.Dense(128, activation='relu'),     tf.keras.layers.Dense(10) ])  # Compile the model with optimizer, loss, and metrics model.compile(     optimizer='adam',     loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),     metrics=['accuracy'] )  # Train the model on training data model.fit(x_train, y_train, epochs=5, batch_size=32)  # Evaluate on test data test_loss, test_acc = model.evaluate(x_test, y_test, verbose=2) print(f"nTest accuracy: {test_acc:.4f}")  # Save the trained model to disk (SavedModel format) model.save("model_saved.keras")

LiteRT: Легковесный вывод на устройстве

LiteRT (Lite Runtime) по сути является облегчённым движком для выполнения выводов, который изначально возник как TensorFlow Lite. LiteRT позволяет запускать заранее обученные модели на устройствах с ограниченными ресурсами (мобильные телефоны, планшеты, устройства IoT, периферийные устройства и микроконтроллеры).

Изначально он был разработан преимущественно для поддержки моделей, созданных с помощью TensorFlow. Однако команда Google AI Edge также расширила LiteRT, чтобы он мог работать с моделями, созданными в других фреймворках. Он использует инструменты конвертации, которые позволяют взять модель из PyTorch, JAX или TensorFlow и преобразовать её в формат FlatBuffers .tflite.

LiteRT: Сильные и слабые стороны

LiteRT разработан для распознавания на устройстве на мобильных, встроенных и периферийных нагрузках. Вот краткий обзор того, в чем он преуспевает — и где проявляются компромиссы:

Сильные стороны LiteRT

  • Ультралёгкое время выполнения – супер маленький исполняемый файл (самая лёгкая сборка ~300 КБ), важно для мобильных/встроенных приложений, которым нужны маленькие бинарные файлы.
  • Поддержка аппаратного ускорения – поддерживает Android NNAPI (DSP/NPU) и iOS Core ML, а также делегат GPU для мобильных GPU.
  • Оптимизация моделей – можно применять пост-тренировочную квантизацию (int8/float16), обрезку, кластеризацию для уменьшения размера и задержки с минимальными потерями точности.
  • Оптимизировано для работы на устройстве – агрессивные оптимизации и делегаты обеспечивают вывод в реальном времени и сниженное энергопотребление.
  • Кроссплатформенная поддержка – поддерживает Android, iOS, Linux (включая Raspberry Pi) и микроконтроллеры (через LiteRT Micro).
  • Офлайн и ориентировано на конфиденциальность — выполняет обработку данных локально на устройстве, избегая сетевой задержки и защищая данные пользователя.

Слабые стороны LiteRT

  • Среда выполнения только для вывода – поддержка общего обучения на устройстве отсутствует (возможна ограниченная поддержка переносного обучения в специализированных сценариях).
  • Пробелы в поддержке операторов — не все операторы TensorFlow/PyTorch поддерживаются нативно; возможно, потребуется переписать модель или создать пользовательские операторы.
  • Трение при конверсии — может потребоваться сегментировать графы и применить квантизацию в некоторых случаях для преобразования модели в формат .tflite.
  • Проблемы отладки — статические модели FlatBuffer сложнее проверять и отлаживать по сравнению с нативными моделями фреймворка.
  • Ограничения памяти/вычислений – крупные модели уровня трансформеров или другие ресурсоемкие модели все еще могут быть слишком медленными для целевых устройств.
  • Возможна необходимость разгрузки сервера — для моделей, которые превышают вычислительные или мемориальные возможности устройства, может потребоваться гибридный клиент-серверный вывод.

LiteRT в ML-пайплайне

Пайплайн можно описать следующим образом: обучаем модель в TensorFlow/PyTorch, затем конвертируем обученную модель в формат .tflite с помощью соответствующего конвертера. Наконец, разворачиваем этот файл .tflite в мобильном приложении или на встроенном устройстве с использованием среды выполнения LiteRT.

Во время конверсии вы можете применять оптимизации, такие как квантизация или обрезка. Эта модель затем будет выполняться как часть вашего программного обеспечения с использованием интерпретатора LiteRT (который можно использовать на многих языках — например, Java/Kotlin на Android, Swift на iOS, C++ для нативного развития или Python для быстрого прототипирования).

В целом, он обладает значительно лучшей производительностью на устройстве по сравнению с использованием полного фреймворка на устройстве. В одном бенчмарке на мобильном устройстве Samsung S21 базовая модель классификации изображений выполняла инференс за 23 мс при ~89 МБ памяти с TensorFlow Lite. Однако та же модель показала ~31 мс (112 МБ) и ~38 мс (126 МБ) с ONNX Runtime и PyTorch Mobile соответственно. Это подчеркивает важность ориентированности LiteRT на низкую задержку и экономное использование памяти при выполнении на мобильных устройствах.

Пример LiteRT – конвертация и использование модели

Этот пример показывает, как преобразовать модель TensorFlow, обученную на Python, в формат LiteRT и выполнить вывод.

# ready MNIST → SavedModel → LiteRT(TFLite) → Inference pipeline import tensorflow as tf import numpy as np  print("TF version:", tf.__version__)  # 1) Load & prep MNIST (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data() # Normalize to [0,1] x_train = (x_train / 255.0).astype("float32") x_test  = (x_test  / 255.0).astype("float32")  # 2) Define & train a simple Keras model model = tf.keras.Sequential([     tf.keras.layers.Input(shape=(28, 28)),     tf.keras.layers.Flatten(),     tf.keras.layers.Dense(128, activation="relu"),     tf.keras.layers.Dense(10)  # logits ])  model.compile(optimizer="adam",               loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),               metrics=["accuracy"])  model.fit(x_train, y_train, epochs=2, batch_size=128, validation_split=0.1, verbose=1)  # Quick test set eval test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0) print(f"TF model test accuracy: {test_acc:.4f}")  # 3) Export a TensorFlow SavedModel (needed for TFLite conversion) # In TF 2.15+, prefer model.export(). If not available, fallback to tf.saved_model.save if hasattr(model, "export"):     model.export("model_saved")            # TF ≥ 2.15 else:     tf.saved_model.save(model, "model_saved")  # Older TF fallback  # 4) Convert to LiteRT/TFLite (FP32 with default optimizations) converter = tf.lite.TFLiteConverter.from_saved_model("model_saved") converter.optimizations = [tf.lite.Optimize.DEFAULT]   # dynamic range quantization if weights permit tflite_model = converter.convert()  with open("model_fp32.tflite", "wb") as f:     f.write(tflite_model) print("Wrote model_fp32.tflite")  # --- OPTIONAL: Full INT8 quantization with representative dataset --- do_full_int8 = True if do_full_int8:     def rep_data():         # Yield a few hundred samples to calibrate ranges         for i in range(500):             # TFLite expects a batch dimension             yield [np.expand_dims(x_train[i], 0)]     converter_int8 = tf.lite.TFLiteConverter.from_saved_model("model_saved")     converter_int8.optimizations = [tf.lite.Optimize.DEFAULT]     converter_int8.representative_dataset = rep_data     # Force int8 I/O where supported     converter_int8.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]     converter_int8.inference_input_type = tf.int8     converter_int8.inference_output_type = tf.int8     try:         tflite_int8 = converter_int8.convert()         with open("model_int8.tflite", "wb") as f:             f.write(tflite_int8)         print("Wrote model_int8.tflite")     except Exception as e:         print("INT8 conversion fell back / failed:", e)         tflite_int8 = None  # 5) Run inference with the TFLite Interpreter (FP32 model) import tensorflow.lite as tflite  def tflite_predict(tflite_path, image_28x28):     interpreter = tflite.Interpreter(model_path=tflite_path)     interpreter.allocate_tensors()     in_details = interpreter.get_input_details()     out_details = interpreter.get_output_details()      inp = image_28x28     # Match dtype & shape expected by the model     if in_details[0]["dtype"] == np.float32:         inp = inp.astype(np.float32)     elif in_details[0]["dtype"] == np.int8:         # Quantized model expects int8; apply quantization params         scale, zero_point = in_details[0]["quantization"]         if scale == 0:             # Safety: if no scale provided (rare), just cast             inp = inp.astype(np.int8)         else:             inp = (inp / scale + zero_point).round().astype(np.int8)      # Add batch dimension     inp = np.expand_dims(inp, 0)      interpreter.set_tensor(in_details[0]["index"], inp)     interpreter.invoke()     out = interpreter.get_tensor(out_details[0]["index"])      # If output is int8, dequantize back to float for softmax/argmax     if out_details[0]["dtype"] == np.int8:         scale, zero_point = out_details[0]["quantization"]         if scale != 0:             out = (out.astype(np.float32) - zero_point) * scale      # Convert logits to probabilities and pick class     probs = tf.nn.softmax(out, axis=-1).numpy()[0]     pred  = int(np.argmax(probs))     conf  = float(probs[pred])     return pred, conf  # Test on a few MNIST samples with FP32 model for idx in [0, 1, 2]:     pred, conf = tflite_predict("model_fp32.tflite", x_test[idx])     print(f"[FP32] Sample {idx}: pred={pred}, conf={conf:.3f}, true={y_test[idx]}")  # If INT8 model exists, test it as well if 'tflite_int8' in locals() and tflite_int8 is not None:     for idx in [0, 1, 2]:         pred, conf = tflite_predict("model_int8.tflite", x_test[idx])         print(f"[INT8] Sample {idx}: pred={pred}, conf={conf:.3f}, true={y_test[idx]}")

Код сначала загружает и нормализует набор данных MNIST, определяет и обучает небольшую полностью связанную сеть, а затем оценивает точность. Обученная модель сохраняется как SavedModel, который используется как вход для TFLiteConverter. Конвертер генерирует две модели: стандартную FP32 LiteRT модель и (опционально) полностью квантизированную модель INT8 с использованием представительного набора данных для калибровки диапазонов. Наконец, код определяет вспомогательную функцию tflite_predict(), которая загружает файл .tflite, подготавливает и при необходимости квантизирует/деквантизирует данные, выполняет вывод и возвращает предсказанную цифру и уверенность. Несколько тестовых примеров проходят через обе модели, FP32 и INT8, чтобы подтвердить корректность развертывания и показать пример выходных данных.

TensorRT: Высокопроизводительное выполнение вывода на графических процессорах NVIDIA

NVIDIA TensorRT — это SDK и среда выполнения для развертывания нейронных сетей с низкой задержкой и высокой пропускной способностью на графических процессорах NVIDIA. TensorRT можно рассматривать как компилятор моделей глубокого обучения: вы предоставляете обученную модель (обычно в формате ONNX или специфическом для фреймворка), и он выполняет серию оптимизаций, чтобы создать оптимизированный движок инференса, который работает на GPU.

Эти оптимизации включают в себя:

  • Слияние слоев — объединение совместимых операций для уменьшения объема передаваемых данных в память и количества запусков ядра.
  • Автонастройка ядра — выбирает самые быстрые CUDA-ядра для заданной формы/аппаратуры.
  • Планирование памяти — оптимизируйте время жизни тензоров и рабочее пространство, чтобы минимизировать копирования и пиковые нагрузки.
  • Сниженная точность — включите поддержку FP16 и INT8 (с калибровкой/QAT) для значительного ускорения и снижения требований к пропускной способности.
  • Динамические формы и кэширование профилей — создавайте профили выполнения для диапазонов форм, чтобы избежать повторной оптимизации во время выполнения.

Окончательный результат — это высоко оптимизированный бинарный файл, который может выполнять прямой проход модели значительно быстрее, чем стандартные реализации.

TensorRT: сильные и слабые стороны

Этот раздел описывает основные функции TensorRT и компромиссы, с которыми вам придется справляться.

Сильные стороны TensorRT

  • Высокая производительность GPU – Использует точность FP16/INT8 и оптимизации, специфичные для архитектуры NVIDIA, чтобы значительно ускорить вывод (выполнение) (часто до ~40× быстрее по сравнению с выполнением на CPU).
  • Эффективность на уровне GPU – может обеспечить задержку в 2–5 раз ниже или выше пропускную способность по сравнению с неоптимизированным TensorFlow/PyTorch на GPU.
  • Пакетная обработка/конкурентность – Разработано для серверов инференса большого масштаба (множество одновременных запросов), где важна пропускная способность.
  • Гибкая интеграция – Может быть интегрирована как автономное решение или в составе Triton Inference Server или ONNX Runtime (используется как поставщик выполнения), либо как часть интеграции с TensorFlow.
  • Широкая поддержка экосистемы – Интегрируется с SDK от NVIDIA, такими как DeepStream (видеоаналитика), Riva (речевые/AI-ассистенты) и другими, чтобы обеспечивать полноценные решения.
  • Поддержка языков – Поддерживается интеграция с API на Python и C++ для включения в пользовательские конвейеры обработки данных.

Слабые стороны TensorRT

  • Зависит от оборудования NVIDIA – TensorRT работает только на графических процессорах NVIDIA. Решения для инференса, построенные на TensorRT, не портируемы на ЦПУ и другие ускорители.
  • Сложность конверсии – модель должна быть преобразована в формат движка TensorRT, что может потребовать переработки для неподдерживаемых операций или написания пользовательских плагинов.
  • Требуемые усилия по настройке – Для достижения оптимальных результатов необходимо вручную настраивать различные параметры, такие как размер рабочей области, режим точности и профили формы ввода.
  • Аппаратно-специфические движки – движки TensorRT оптимизированы для конкретного семейства и архитектуры GPU. Вы можете пересобрать их для другого оборудования при переносе вашего приложения (например, с GPU центров обработки данных на Edge-устройство Jetson).
  • Перестройка при обновлениях модели – Любое изменение модели требует повторного преобразования и повторной оптимизации для поддержания производительности.
  • Сложность развертывания – усилия по разработке и сопровождению выше по сравнению с интегрированными решениями, такими как TF Serving или только ONNX Runtime.

TensorRT в цепочке обработки: только вывод

TensorRT предназначен только для времени вывода. Например, вы обучаете модель с помощью PyTorch или TensorFlow на GPU, затем экспортируете её в ONNX (TensorRT поддерживает ONNX как стандартный формат ввода). После этого вы можете использовать API или утилиты TensorRT для развертывания этой модели ONNX, создания движка (с калибровкой на примерах данных, если используется квантование INT8), и запускать этот движок в вашем приложении.

Теперь загрузите движок в серверное приложение на C++ или в привязки Python (если у вас небольшая установка). Если вы работаете с сервером высокой пропускной способности (например, обработка миллионов запросов в производстве), очень часто движки TensorRT запускаются в NVIDIA Triton Inference Server, который управляет несколькими моделями и параллельной обработкой.

Пример TensorRT – Конвертация модели ONNX в движок TensorRT

Ниже приведен упрощенный (псевдокод) пример использования Python API TensorRT для построения механизма из модели ONNX и выполнения инференса. Это дает представление о общем рабочем процессе, без деталей API.

import tensorrt as trt  onnx_file = "model.onnx" engine_file = "model.plan"  # TensorRT engine file  # Set up TensorRT logger and builder logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger) network = builder.create_network(flags=1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger)  # Parse the ONNX model to populate the TensorRT network with open(onnx_file, "rb") as f:     parser.parse(f.read()) # (In practice, check parser.error for unsupported ops here.)  # Configure builder builder.max_batch_size = 1 config = builder.create_builder_config() config.max_workspace_size = 1 << 30  # 1GB workspace for optimization config.flags |= trt.BuilderFlag.FP16  # enable FP16 precision if supported  # Build the TensorRT engine engine = builder.build_engine(network, config) with open(engine_file, "wb") as f:     f.write(engine.serialize())  # save the engine to file  # Use the engine for inference runtime = trt.Runtime(logger) with open(engine_file, "rb") as f:     engine_bytes = f.read() engine = runtime.deserialize_cuda_engine(engine_bytes) context = engine.create_execution_context()  # Assuming a single input and single output for simplicity input_shape = engine.get_binding_shape(0) output_shape = engine.get_binding_shape(1) # Allocate device memory for inputs and outputs (using PyCUDA or similar) # ... (omitted for brevity) # Execute inference context.execute_v2(bindings=[d_input_ptr, d_output_ptr]) # Copy results from device memory to host and use the output

Приведённый выше код описывает шаги, необходимые для преобразования модели ONNX в движок TensorRT. Сначала мы создаём Builder и OnnxParser, которые используются для чтения графа модели. Далее задаются некоторые параметры конфигурации билдера (такие как размер рабочей области и включение поддержки FP16). Вызов build_engine выполняет оптимизацию TensorRT и выводит движок (сериализованный в “model.plan”). Позже мы можем десериализовать движок и создать контекст выполнения, который можно использовать для запуска инференса.

Давайте рассмотрим шаги, необходимые для использования TensorRT для вывода:

  • Разбор модели: Загрузите и интерпретируйте структуру модели с помощью IParser TensorRT или других соответствующих инструментов.
  • Создание оптимизированного движка: Создайте оптимизированный движок на основе распарсенной модели, который будет специфичен для целевого оборудования. Обычно это включает оптимизации, такие как слияние слоев и калибровка точности.
  • Запуск вывода: Выполните вывод с использованием оптимизированного движка. Этот шаг включает управление памятью устройства для входных и выходных данных и вызов метода execute_v2 на контексте выполнения.

В реальных приложениях также необходимо обрабатывать динамические формы или несколько привязок, но это требует гораздо более глубокого погружения в код TensorRT. На практике большинство разработчиков используют высокоуровневую обертку или ONNX Runtime с бэкендом TensorRT, чтобы избежать написания низкоуровневого кода самостоятельно.

ONNX: Взаимодействие моделей и кроссплатформенное развертывание

ONNX (Open Neural Network Exchange) не является фреймворком для обучения; это открытый формат для представления моделей машинного обучения, который предоставляет среду выполнения (ONNX Runtime) для выполнения моделей.

Вы можете обучить модель в одной среде (например, PyTorch или TensorFlow), экспортировать её в формат ONNX (вычислительный граф с стандартными операциями), а затем запускать её с использованием другого инструмента или даже другой аппаратной платформы. Разделение сред разработки и времени выполнения очень мощное в производственной среде, где вы можете выбрать лучшую среду для разработки, а затем подобрать лучшее время выполнения для развёртывания.

ONNX: Сильные и слабые стороны

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

Сильные стороны (ONNX/ONNX Runtime)

  • Гибкость и совместимость: Обучайте в одной среде (например, PyTorch), развертывайте в другой (например, ORT на C++).
  • Лёгкая рабочая среда, ориентированная на вывод: размер установки меньше, чем у полноценных фреймворков. Нет расходов на обучение.
  • Кроссплатформенность: Windows, Linux, macOS и мобильные устройства (ORT Mobile с ограниченным функционалом).
  • Высокая производительность: На уровне или лучше, чем у родного фреймворка при инференсе, особенно на CPU; при некоторых запусках на GPU с batch-1 ORT может работать быстрее после оптимизации графа (например, ~24,2 мс против ~30,4 мс на ResNet-50).
  • Доказано в масштабах: Сообщается, что превосходит TorchScript в некоторых производственных нагрузках Microsoft.
  • Здоровая экосистема: Здоровая экосистема конвертеров, зоопарков моделей и инструментов (Netron, onnxoptimizer и др.).

Слабые стороны (ONNX/ONNX Runtime)

  • Сложность конвертации: Экспорт может завершиться сбоем или привести к снижению производительности, если модель содержит операции, которые еще не стандартизированы/не поддерживаются.
  • Пробелы в пользовательских операциях: Модели с нестандартными слоями могут требовать резервных вариантов или пользовательских плагинов. Они могут превращаться в «черные ящики».
  • Сложнее отлаживать: Статические графы ONNX трудно отлаживать, так как они отделены от исходного фреймворка, исходного кода и инструментов.
  • Меньше «всё включено»: Пред- и постобработка, а также другие «связующие» элементы конвейера часто выполняются вне ORT.
  • Накладные расходы рабочего процесса: Добавляет дополнительный этап экспорта/проверки в рабочий процесс, который нужно поддерживать.

ONNX как слой передачи: обучай где угодно, развёртывай повсюду

ONNX часто находится посредине конвейера. Распространённый сценарий заключается в том, чтобы обучить модель в PyTorch, а затем использовать torch.onnx.export для экспорта модели в model.onnx. После этого мы можем взять эту ONNX-модель и развернуть её в производственном сервисе (с использованием ONNX Runtime, который может быть написан на C++ для повышения эффективности, или на Python, если это подходит).

Давайте предположим, что вы работаете с моделью TensorFlow, но хотели бы использовать TensorRT без интеграции с TensorFlow. В этом случае вы можете конвертировать модель TF в ONNX, а затем передать её TensorRT (так как TensorRT также имеет встроенную поддержку ONNX).

Аналогично, ONNX используется в процессах сжатия и квантизации моделей. Например, вы можете экспортировать модель в формат ONNX, а затем выполнить посттренировочную квантизацию с помощью инструментов, предоставленных ONNX Runtime.

Взаимодействие и типичные рабочие процессы

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

Трубопровод Цель/Окружение Ключевые шаги
PyTorch → ONNX → TensorRT (развертывание на GPU) Серверы/устройства на базе NVIDIA GPU с поддержкой CUDA Обучите в PyTorch. Экспортируйте в ONNX (выберите opset, упростите граф). Создайте движок TensorRT (FP16/INT8, профиль). Разверните движок и обрабатывайте запросы.
TensorFlow → LiteRT (Мобильное развертывание) Android и iOS (на устройстве) Обучайте в TF/Keras. Конвертируйте с помощью LiteRT (TFLite) конвертера в .tflite. Включите в приложение; активируйте делегаты. (Необязательно) QAT через оптимизацию моделей TF
PyTorch → LiteRT (Прямой или через ONNX) Android и iOS (на устройстве) Обучайте в PyTorch. Конвертируйте напрямую в .tflite или через ONNX + конвертер TF. Интегрируйте в мобильные приложения.
PyTorch → ONNX → ONNX Runtime (CPU/GPU) Windows, Linux, macOS, мобильные Обучайте в предпочитаемой среде. Экспортируйте в ONNX. Запускайте с ONNX Runtime (выбирайте поставщика в зависимости от платформы).
TensorFlow → TensorRT (TF-TRT или ONNX) Серверы с графическими процессорами NVIDIA Вариант A: TF-TRT (части графа заменены на движки TRT; совместимо с TF Serving). Вариант B: Экспорт в ONNX → прямое создание движка TRT

PyTorch и TensorFlow — это «фронтенды», используемые для построения моделей; ONNX — это общий «промежуточный» формат для переноса моделей из одного фреймворка в другой; TensorRT и LiteRT — это «конечные точки», каждая оптимизирована для конкретного оборудования (GPU и устройств на периферии соответственно).

Часто задаваемые вопросы

  • PyTorch против TensorFlow — когда стоит выбирать каждый из них? PyTorch отлично подходит для быстрой итерации исследований и удобной отладки в стиле Python; TensorFlow лучше использовать для полноценных производственных ML-пайплайнов (TFX, TF-Serving, TPU) и более плавной работы в корпоративной среде.

  • Что такое LiteRT (ранее TFLite) и когда его использовать? LiteRT — это лёгкая среда выполнения для вывода на устройстве, предназначенная для мобильных и периферийных устройств. Обучайте свою модель с помощью TensorFlow или PyTorch, конвертируйте в .tflite и запускайте с аппаратными делегатами (NNAPI, Core ML, GPU) для быстрого и энергоэффективного вывода.

  • Как работают ONNX и TensorRT вместе? Экспортируйте вашу обученную модель в ONNX, затем используйте TensorRT для компиляции её в высоко оптимизированный движок для графических процессоров NVIDIA (FP16/INT8, слияние ядер). ONNX — это мост; TensorRT — это турбонаддув для GPU.

Заключение

Выбирайте инструменты исходя из того, где вы запускаете и как масштабируете: PyTorch или TensorFlow для итерации и обучения; ONNX для разделения обучения и обслуживания; TensorRT для максимальной производительности на GPU NVIDIA; и LiteRT для небольшого, низколатентного инференса на устройстве. Большинство успешных стеков используют несколько фреймворков (например, PyTorch → ONNX → TensorRT для обслуживания на GPU или TensorFlow → LiteRT для мобильных устройств) с единым экспортированным артефактом, который можно тестировать и внедрять.

Практический способ добиться этого — с помощью платформы Linux-Console.net Gradient AI. Запускайте управляемые GPU-блокноты, обучайте модели и развёртывайте ускоренные конечные точки без управления инфраструктурой. Это позволит вам использовать PyTorch/TensorFlow с ONNX, TensorRT или LiteRT, всё в одном упрощённом рабочем процессе.

Ссылки и ресурсы

  • Сравнительное исследование PyTorch и TensorFlow для глубокого обучения: удобство использования, производительность и компромиссы при внедрении
  • Использование интеграции TensorFlow с TensorRT для низколатентного вывода
  • Обзор LiteRT
  • TensorFlow Lite против PyTorch Mobile для машинного обучения на устройстве
  • Edge AI: TensorFlow Lite против ONNX Runtime против PyTorch Mobile
  • ONNX, ONNX Runtime и TensorRT
  • Скорость ONNX против PyTorch: Подробное сравнение производительности
  • Масштабирование инференса PyTorch: Обслуживание миллиардов ежедневных NLP-инференсов с помощью ONNX Runtime

Спасибо, что учитесь вместе с сообществом Linux-Console.net.

Комментарии

Добавить комментарий

Ваш адрес email не будет опубликован. Обязательные поля помечены *