TensorRT onnx转engine报错:
[TRT] [E] 2: [ltWrapper.cpp::nvinfer1::rt::CublasLtWrapper::setupHeuristic::334] Error Code 2: Internal Error
(Assertion cublasStatus == CUBLAS_STATUS_SUCCESS failed. )
环境说明:
tensorrt:8.X
pytorch1.7
cuda:10.2
cudnn:8.2.1
win10
Nvidia 1650 4G
python 3.7
我用yolov5转onnx再转engine没有问题。然后Resnet onnx转engine出现了错误,根据报错说是CUDA10.2需要打两个补丁,这两个补丁我也打了还是未解决。
通过Debug进行分析,可以通过onnx解析器构建network。并且可以获得正常的输入和输出结点名字以及shape,这个报错是发生在写入resnet.engine文件的时候出错的。
还有一点,如果我直接用tensorrt下bin目录下的trtexec.exe 将onnx转engine是可以成功的。所以不知道问题到底出现在了哪里。
下面是我的onnx转engine代码。
分析的的时候时发现**engineString = builder.build_engine(network, config)**返回为None,所以无法写入,但我network是已经成功构建好的呀,所以很奇怪。
def export_engine(model, im, file, half, workspace=4, verbose=False):
# 首先判断一下im是不是在GPU上
assert im.device.type != 'cpu', 'export running on CPU but must be on GPU, i.e.'
import tensorrt as trt
# 判断trt版本
if trt.__version__[0] == '8':
check_version(trt.__version__, '8.0.0', hard=True) # require tensorrt>=8.0.0
# 先转onnx
onnx = file.with_suffix('.onnx') # 获取权重名
LOGGER.info(f'\n{prefix} starting export with TensorRT {trt.__version__}...')
assert onnx.exists(), f'failed to export ONNX file: {onnx}'
f = file.with_suffix('.engine') # TensorRT engine file
# 记录trt转engine日志
logger = trt.Logger(trt.Logger.INFO)
if verbose:
logger.min_severity = trt.Logger.Severity.VERBOSE
# 1.builder构造,记录日志
builder = trt.Builder(logger)
# 2.builder.config建立
config = builder.create_builder_config()
# 3.workspace 构建期间可用显存 workspace * 1 << 30 表示将workspace * 1 二进制左移30位后的10进制
config.max_workspace_size = workspace * 1 << 30
# 4.定义Network
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger) # 使用ONNX解析器导入model
success = parser.parse_from_file(str(onnx))
for idx in range(parser.num_errors):
LOGGER.info(parser.num_errors(idx))
if not success:
raise RuntimeError(f'failed to load ONNX file: {onnx}')
# 5.加载onnx解析器
# if not parser.parse_from_file(str(onnx)):
# raise RuntimeError(f'failed to load ONNX file: {onnx}')
# 6.获得网络输入输出
inputs = [network.get_input(i) for i in range(network.num_inputs)]
outputs = [network.get_output(i) for i in range(network.num_outputs)]
# 下面的只是在log中打印input和output 的name和shape以及数据类型
for inp in inputs:
LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}')
for out in outputs:
LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}')
LOGGER.info(f'{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}')
# 判断是否支持FP16推理
if builder.platform_has_fast_fp16 and half:
config.set_flag(trt.BuilderFlag.FP16)
# build engine 文件的写入 这里的f是前面定义的engine文件
engineString = builder.build_engine(network, config)
if engineString == None:
LOGGER.info("Failed getting serialized engine!")
return
LOGGER.info("Succeeded getting serialized engine!")
with open(f, 'wb') as t:
t.write(engineString)
LOGGER.info("Succeeded saving .plan file!")