Description
I met the problem only with 1080ti device. But with 2080ti, the conversion works fine. The error looks like this:
Traceback (most recent call last):
File "convert_onnx_to_trt.py", line 94, in <module>
main()
File "convert_onnx_to_trt.py", line 90, in main
_ = build_engine(onnx_file_path, engine_file_path, verbose)
File "convert_onnx_to_trt.py", line 49, in build_engine
engine = builder.build_cuda_engine(network)
RuntimeError: Driver error:
Environment
TensorRT Version: 7.1.3.4
GPU Type: 1080ti
Nvidia Driver Version: 455.45.01
CUDA Version: 11.1
CUDNN Version: 8.0.4
Operating System + Version: docker Ubuntu 18.04.5 LTS on Ubuntu 16.04.4
Python Version (if applicable): Python 3.6.5 :: Anaconda, Inc.
TensorFlow Version (if applicable):
PyTorch Version (if applicable):
Baremetal or Container (if container which image + tag):
Relevant Files
Steps To Reproduce
run convert code:
from __future__ import print_function
import os
import argparse
import tensorrt as trt
import pycuda.autoinit
import pycuda.driver as cuda
ctx = cuda.Device(0).make_context()
EXPLICIT_BATCH = []
if trt.__version__[0] >= '7':
EXPLICIT_BATCH.append(
1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
def build_engine(onnx_file_path, engine_file_path, verbose=False):
"""Takes an ONNX file and creates a TensorRT engine."""
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE) if verbose else trt.Logger()
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(*EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = 1 << 28
builder.max_batch_size = 1
builder.fp16_mode = True
#builder.strict_type_constraints = True
# Parse model file
if not os.path.exists(onnx_file_path):
print('ONNX file {} not found, please run yolov3_to_onnx.py first to generate it.'.format(onnx_file_path))
exit(0)
print('Loading ONNX file from path {}...'.format(onnx_file_path))
with open(onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
if not parser.parse(model.read()):
print('ERROR: Failed to parse the ONNX file.')
for error in range(parser.num_errors):
print(parser.get_error(error))
return None
if trt.__version__[0] >= '7':
# The actual yolov3.onnx is generated with batch size 64.
# Reshape input to batch size 1
shape = list(network.get_input(0).shape)
shape[0] = 1
# network.get_input(0).shape = shape
network.get_input(0).shape = ModelData.INPUT_SHAPE
print('Completed parsing of ONNX file')
print('Building an engine; this may take a while...')
engine = builder.build_cuda_engine(network)
print('Completed creating engine')
print('engine',engine)
print('engine_file_path:',engine_file_path)
with open(engine_file_path, 'wb') as f:
f.write(engine.serialize())
return engine
class ModelData(object):
# MODEL_FILE = 'onnx.save/fcn_hrnet_stick_3.0.onnx'
MODEL_FILE = 'onnx.save/fcn_hrnet_stick_3.0_fix_size.onnx'
INPUT_NAME = 'x'
INPUT_SHAPE = (1,1,416,416)
OUTPUT_NAME = 'save_infer_model/scale_0.tmp_0'
def main():
"""
Discription:
Create a TensorRT engine for ONNX-based Model.
Comand e.g.:
python convert_onnx_to_trt.py -v --model_path ./onnx.save/fcn_hrnet_v1.2_gray.onnx --input_shape 512 --save_path ./trt.save/fcn_hrnet_test.trt
"""
parser = argparse.ArgumentParser()
parser.add_argument('-v', '--verbose', action='store_true',
help='enable verbose output (for debugging)')
parser.add_argument('--model_path', type=str, required=True, help='model path')
parser.add_argument('--input_shape', type=int, required=True, help='input size')
parser.add_argument('--save_path', type=str, required=True, help='output trt file name')
args = parser.parse_args()
# onnx_file_path = '%s.onnx' % args.model
# engine_file_path = '%s.trt' % args.model
ModelData.MODEL_FILE = args.model_path
ModelData.INPUT_SHAPE = (1, 1, args.input_shape, args.input_shape) # N C H W
onnx_file_path = ModelData.MODEL_FILE
engine_file_path = args.save_path# 'c-'
verbose = args.verbose
_ = build_engine(onnx_file_path, engine_file_path, verbose)
if __name__ == '__main__':
main()