Вокруг так много фреймворков для инференса нейронных сетей, что сложно понять, какой именно подойдет тебе лучше всего. Я решил, что реализую одну и ту же задачу на нескольких разных технологиях. Так и родился этот репозиторий.
Задача
Так как целью проекта являлся инференс, то было решено не обучать что-то с нуля, а взять уже готовые модели.
Хотелось выбрать такую задачу, чтобы она отвечала следующим требованиям:
Была не слишком простой (состояла из одной модели), но и не слишком сложной;
Содержала нейросети, обученные под разные задачи (например, сетка под детекцию и под сегментацию);
Имела реализацию на pytorch (потому что я его знаю).
Выбор пал на задачу распознавания российских номеров автомобилей. Модели были взяты из этого репозитория.
Пайлпайн распознавания номеров следующий:
Детекция номеров с помощью Yolov5;
Вырезанные номера прогоняются через Spatial transformer для выравнивания;
Текст номера распознается с LPR-net.
Останавливаться на том, как работают сетки, не буду (тут есть ссылки на сетки, кому интересно).
Нейросети
Так как планируется несколько реализаций инференса, код с ними был вынесен в отдельный пакет, из которого мпортируются сетки.
В качестве менеджера зависимостей был выбран Poetry. Хочется, чтобы на данный проект можно было опираться при реализации своего инференса. Для этого нужно, чтобы его можно было запустить и через месяц, и через год. Poetry как раз фиксирует все зависимости (как основные, так и зависимости зависимостей) в специальном файле poetry.lock, что даст воспроизводимость.
Веса нейросетей были сохранены прямо в репозитории, потому что:
Неохота их куда-то загружать, а потом следить, что ты их случайно не удалил;
Так легче воспроизводить (не нужно что-то качать, а потом думать, под каким именем в какую папку положить);
Они не так уж много весят.
Как их использовать, можно понять из тестов.
Фреймворк
Для инференса используется Nvidia Triton Inference Server (далее Triton).
Я бы сказал, что Triton нацелен на максимальную утилизацию вашего железа:
Позволяет передавать данные между моделями без копирования;
Не грузит модели, которые не используются;
Использует shared memory, чтобы получать данные и отдавать ответы;
И другое.
Основные недостатки, на мой взгляд:
Работает только на видеокартах Nvidia;
Спроектирован под работу на одной машине. Triton не предусматривает, что вы запустите часть моделей на одной физической машине, а часть на другой;
Документация состоит из разбросанных по GitHub-у ридми, часть которых дублируется в сайт.
Конвертируем модели
Pytorch модели
Для использования Triton вам необходимо определить model_repository. Он представляет собой папку с папками под каждую модель. Далее вы сможете объединить модели в некоторый граф вычислений, описав передачу данных между ними.
Папка модели должна содержать конфигурацию модели и её веса.
Triton поддерживает большое число фреймоворков (бекендов) для обучения сеток. Нам понадобятся Pytorch и Python.
Pytorch бекенд позволяет ранить TorchScript модельки. В него у меня переведены LPR-net и STN.
Рассмотрим конвертацию на примере STN. Для этого:
1. Конвертируем в TorchScript
dummy_input = torch.randn(1, 3, 24, 94, device=device)
model = load_stn(settings.STN.WEIGHTS, device)
traced_model = torch.jit.trace(model, [dummy_input])2. Сохраняем в папку с версией модели. Triton хочет, чтобы все ваши модели лежали в отдельной папке с номером версии. Далее вы можете слать запросы на определенную версию. Также с помощью этого механизма можно заменять версии моделек прямо на рабочем проде.
model_path = (
Path(__file__).resolve().parents[1]
/ "model_repository"
/ "stn"
/ "1"
/ "model.pt"
)
traced_model.save(model_path)3. Задаем конфигурацию. Основное, что нужно определить - бекенд, батчайсз (чтобы Triton мог процесить батчами), входы и выходы в модель. В нашем случае получаем:
backend: "pytorch"
max_batch_size: 32
input [
{
name: "input__0"
data_type: TYPE_FP32
dims: [ 3, 24, 94 ]
}
]
output [
{
name: "output__0"
data_type: TYPE_FP32
dims: [ 3, 24, 94 ]
}
]Hidden text
Обратите внимание на нижние подчеркивания в названия входов и выходов. Это специальный формат для Pytorch бекенда.
По итогу получилась такая папка с моделью
Python модели
Если по какой-то причине конвертация не проходит или есть что-то помимо модели, есть Python бекенд. В нашем случае внутри yolo встроена дополнительная логика, если послать в нее numpy или torch тензор. Чтобы не дописывать дополнительный пре- и постпроцессинг был использован Python бекенд.
Структура папки такая же, как и в предыдущем случае, но вместо весов у нас файл model.py. В нем должен быть реализован класс TritonPythonModel. В нем при инициализации создается модель:
def initialize(self, args):
self.model_config = json.loads(args["model_config"])
self.model = load_yolo(
settings.YOLO.WEIGHTS, settings.YOLO.CONFIDENCE, torch.device("cuda")
)Для каждого поступившего в модель запроса получаются предикты и формируются ответы:
def execute(
self, requests: List[pb_utils.InferenceRequest]
) -> List[pb_utils.InferenceRequest]:
responses = []
for request in requests:
image = pb_utils.get_input_tensor_by_name(request, "input__0").as_numpy()
image = prepare_detection_input(image)
detection = self.model(image, size=settings.YOLO.PREDICT_SIZE)
df_results = detection.pandas().xyxy[0]
img_plates = prepare_recognition_input(
df_results, image, return_torch=False
)
out_tensor_0 = pb_utils.Tensor(
"output__0", img_plates.astype(self.output0_dtype)
)
out_tensor_1 = pb_utils.Tensor(
"output__1",
df_results[["xmin", "ymin", "xmax", "ymax"]]
.to_numpy()
.astype(self.output1_dtype),
)
inference_response = pb_utils.InferenceResponse(
output_tensors=[out_tensor_0, out_tensor_1]
)
responses.append(inference_response)
return responsesЧтобы это заработало, необходимо доставить пакетов в интерпретатор, который Triton использует для прогона ваших моделей на бекенде Python.
В Docker образах Triton используется Python 3.8. Если у вас другая версия, то нужно собирать под себя специально. Я пошёл проще, и использовал 3.8
Интересные статьи
Интересные статьи
