TensorRT不支持的操作\u填充

2024-10-03 02:46:51 发布

您现在位置:Python中文网/ 问答频道 /正文

我试图通过TensorRT引擎将resnet从.pb(tensorflow)运行到.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] ERROR: UffParser: Validator error: reshape_4/zeros: Unsupported operation _Fill
Building TensorRT engine, this may take a few minutes...
[TensorRT] ERROR: Network must have at least one output
[TensorRT] ERROR: Network validation failed.

问题:

在tensorRT中添加不支持的操作需要做什么


Tags: pathimportparsermodelasbuildererrornetwork