Плагин Tensorrt и анализатор кафе в Python

Я новичок в Tensorrt, и я не так хорошо знаком с языком Си. Могу я спросить, есть ли какой-нибудь пример для импорта caffe modell(caffeparser) и в то же время использовать плагин с python. Пример библиотеки плагинов: " https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/c_api/_nv_infer_plugin_8h_source.html".

Я видел пример, делающий что-то вроде ниже. Нужно ли модифицировать класс плагина? или это уже было сделано с помощью плагина Python API?

import tensorrt
import tensorrtplugins
from tensorrt.plugins import _nv_infer_plugin_bindings as nvinferplugin
from tensorrt.parsers import caffeparser

plugin_factory = tensorrtplugins.FullyConnectedPluginFactory()
parser = caffeparser.create_caffe_parser()
parser.set_plugin_factory(plugin_factory)

engine = trt.utils.caffe_to_trt_engine(G_LOGGER,
MODEL_PROTOTXT,
CAFFE_MODEL,
1,
1 << 20,
OUTPUT_LAYERS,
trt.infer.DataType.FLOAT,
plugin_factory
)

PS: я пытаюсь конвертировать YOLO2 в формат Tensorrt. Поэтому некоторые слои (например, kYOLOREORG и kPRELU) могут поддерживаться только плагином.

Еще один способ сделать это - добавить плагин во время построения сети, методом network.add_plugin_ext(). Однако я не уверен, как указать предыдущий слой, который будет импортирован позже.

Большое спасибо за ваш ответ. Ваша помощь будет высоко ценится!

0 ответов

Другие вопросы по тегам