Развертывание сети семантической сегментации (U-Net) с TensorRT (без поддержки повышающей дискретизации)

Я пытаюсь развернуть обученную U-Net с TensorRT. Модель была обучена с использованием Keras (с Tensorflow в качестве бэкэнда). Код очень похож на этот: https://github.com/zhixuhao/unet/blob/master/model.py

Когда я преобразовал модель в формат UFF, используя такой код:

import uff
import os
uff_fname = os.path.join("./models/", "model_" + idx + ".uff")
uff_model = uff.from_tensorflow_frozen_model(
    frozen_file = os.path.join('./models', trt_fname), output_nodes = output_names, 
    output_filename = uff_fname
)

Я получу следующее предупреждение:

Warning: No conversion function registered for layer: ResizeNearestNeighbor yet.
Converting up_sampling2d_32_12/ResizeNearestNeighbor as custom op: ResizeNearestNeighbor
Warning: No conversion function registered for layer: DataFormatVecPermute yet.
Converting up_sampling2d_32_12/Shape-0-0-VecPermuteNCHWToNHWC-LayoutOptimizer as custom op: DataFormatVecPermute

Я попытался избежать этого, заменив слой повышающей дискретизации с повышающей дискретизацией (билинейная интерполяция) и транспонируя свертку. Но конвертер выдаст мне подобные ошибки. Я проверил https://docs.nvidia.com/deeplearning/sdk/tensorrt-support-matrix/index.html, и казалось, что все эти операции еще не поддерживаются.

Мне интересно, есть ли какое-либо решение этой проблемы? Есть ли другой формат/фреймворк, который TensorRT любит и поддерживает повышающую дискретизацию? Или его можно заменить другими поддерживаемыми операциями?

Я также видел где-то, что можно добавить настраиваемые операции, чтобы заменить те, которые не поддерживаются для TensorRT. Хотя я не совсем уверен, каким будет рабочий процесс. Также было бы очень полезно, если бы кто-то мог указать на пример пользовательских слоев.

Заранее спасибо!

Ответ 1

Предупреждения заключаются в том, что эти операции еще не поддерживаются TensorRT, как вы уже упоминали. К сожалению, нет простого способа исправить это. Вы должны либо изменить график (даже после обучения), чтобы использовать только операцию, поддерживаемую комбинацией; или напишите эти операции самостоятельно как пользовательский слой.

Однако в C++ есть лучший способ выполнить логический вывод на других устройствах. Вы можете использовать TensorFlow вместе с TensorRT вместе. TensorRT проанализирует график на предмет поддерживаемых операций и преобразует их в узлы TensorRT, а оставшаяся часть графика будет обрабатываться TensorFlow как обычно. Больше информации здесь. Это решение намного быстрее, чем переписывание операций самостоятельно. Единственная сложная часть - это построить TensorFlow из исходных tensorflow_cc на целевом устройстве и сгенерировать динамическую библиотеку tensorflow_cc. В последнее время появилось много руководств и поддержка портов TensorFlow для различных архитектур, например, ARM.

Ответ 2

Эй, я сделал что-то похожее, я бы сказал, что лучший способ решить эту проблему - экспортировать вашу модель в .onnx с такой же, как эта, если вы проверите матрицу поддержки для onnx поддерживается upsample:

Затем вы можете использовать https://github.com/onnx/onnx-tensorrt для преобразования onnx-модели в tenorrt, я хочу преобразовать сеть, которую я обучил Pytorch, и которая имела повышенный уровень. Репозиторий для onnx-tensorrt немного более активен, и если вы откроете вкладку pr, вы можете проверить, что другие люди пишут пользовательские слои и оттуда разветвляются.

Ответ 3

Обновление от 28.09.2009

Nvidia выпустила TensorRT 6.0.1 около двух недель назад и добавила новый API под названием "IResizeLayer". Этот уровень поддерживает "Ближайшую" интерполяцию и, таким образом, может использоваться для реализации повышающей дискретизации. Больше не нужно использовать пользовательские слои/плагины!

Оригинальный ответ:

Спасибо за все ответы и предложения, размещенные здесь!

В итоге мы напрямую внедрили сеть в API TensorRT C++ и загрузили весовые коэффициенты из файла модели .h5. У нас пока нет времени профилировать и полировать решение, но, похоже, вывод работает в соответствии с тестовыми изображениями, которые мы ввели.

Вот рабочий процесс, который мы приняли:

Шаг 1: закодируйте слой повышенной дискретизации.

В нашей модели U-Net все уровни повышающей дискретизации имеют коэффициент масштабирования (2, 2), и все они используют интерполяцию ResizeNearestNeighbor. По существу, значение пикселя в точке (x, y) в исходном тензоре будет равно четырем пикселям: (2x, 2y), (2x + 1, 2y), (2x, 2y + 1) и (2x + 1, 2y + 1) ) в новом тензоре. Это может быть легко закодировано в функцию ядра CUDA.

После того как мы получили ядро с повышающей дискретизацией, нам нужно обернуть его TensorRT API, в частности, класс IPluginV2Ext. Справочник разработчика содержит некоторые описания того, какие функции необходимо реализовать. Я бы сказал, что enqueue() - самая важная функция, потому что там выполняется ядро CUDA.

Есть также примеры в папке TensorRT Samples. Для моей версии эти ресурсы полезны:

Github: Leaky Relu как пользовательский слой
TensorRT-5.1.2.2/образцы/sampleUffSSD
TensorRT-5.1.2.2/образцы/sampleSSD

Шаг 2. Кодирование остальной части сети с помощью TensorRT API

Остальная часть сети должна быть довольно простой. Просто найдите функцию вызова addxxxLayer из определения сети TensorRT.

Имейте в виду одну вещь: в зависимости от того, какую версию TRT вы используете, способ добавления отступов может отличаться. Я думаю, что новейшая версия (5.1.5) позволяет разработчикам добавлять параметры в addConvolution(), чтобы можно было выбрать правильный режим заполнения.

Моя модель была обучена с использованием Keras, режим заполнения по умолчанию состоит в том, что правый и нижний получают больше заполнения, если общее количество заполнения не является четным. Проверьте эту ссылку Qaru для получения подробной информации. В 5.1.5 есть режим mode, который представляет эту схему заполнения.

Если вы используете более старую версию (5.1.2.2), вам нужно будет добавить заполнение как отдельный слой перед слоем свертки, который имеет два параметра: предварительное заполнение и последующее заполнение.

Кроме того, все вещи NCHW в TensorRT

Полезный пример:

TensorRT-5.1.2.2/образцы/sampleMNISTAP

Шаг 3. Загрузите веса

TensorRT требует весов в формате [out_c, in_c, filter_h, filter_w], который упоминается в архивированной документации. Вес Keras имеет формат [filter_h, filter_w, c_in, c_out].

Мы получили файл чистого веса, вызвав model.save_weights('weight.h5') в Python. Затем мы можем считать веса в массив Numpy с помощью h5py, выполнить транспонирование и сохранить транспонированные веса в новый файл. Мы также выяснили имя группы и набора данных, используя h5py. Эта информация использовалась при загрузке весов в код C++ с использованием HDF5 C++ API.

Мы сравнивали выходной слой за слоем между кодом C++ и кодом Python. Для нашей сети U-Net все карты активации одинаковы, возможно, до третьего блока (после 2 пулов). После этого между значениями пикселей есть небольшая разница. Абсолютная ошибка в процентах составляет 10 ^ -8, поэтому мы не думаем, что это плохо. Мы все еще находимся в процессе доработки реализации C++.

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