Я пытаюсь запустить re snet из .pb (тензор потока) с помощью двигателя TensorRT в .trt. Я преобразовал свой .pb в .uff и теперь пытаюсь загрузить его в движок с помощью этого кода:
import tensorrt.legacy as trt
import tensorflow as tf
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
from tensorrt.legacy.parsers import uffparser
import graphsurgeon as gs
# Build TensorRT engine
uff_model_path = "model/resnet_model_v1.uff"
engine_path = "model/resnet_model_v1.engine"
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
trt_runtime = trt.Runtime(TRT_LOGGER)
with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
builder.max_workspace_size = 1 << 30
builder.fp16_mode = True
builder.max_batch_size = 1
parser.register_input("input_image", (3, 150, 150))
parser.register_output("embedding_layer/MatMul")
parser.parse(uff_model_path, network)
print("Building TensorRT engine, this may take a few minutes...")
trt_engine = builder.build_cuda_engine(network)
Но он падает с этой ошибкой:
[TensorRT] ОШИБКА: UffParser: ошибка валидатора: reshape_4 / нули: неподдерживаемая операция _Fill
Построение движка TensorRT, это может занять несколько минут ...
[TensorRT] ОШИБКА: Сеть должна иметь хотя бы один выход
[TensorRT] ОШИБКА : Проверка сети не удалась.
ВОПРОС
Что мне нужно сделать с добавлением неподдерживаемой операции в tenorRT?