Ubuntu20.04配置TensorRT

本文主要是介绍Ubuntu20.04配置TensorRT，对大家解决编程问题具有一定的参考价值，需要的程序猿们随着小编来一起学习吧！

linux20.04 + CUDA11.3 + cudnn8.2 + TensorRT8.0.1.6

nvidia-smi显示的CUDA版本和nvcc -V得到的CUDA版本会不一致，nvidia-smi显示的是支持的最高的

1.安装CUDA，从官网下载，可下在run的和deb的，按照提示的命令安装，安装完成后/usr/local/cuda*路径下
2.安装cudnn，从官网下载，下载tar版本的，解压后有include和lib64文件夹，执行如下命令安装

sudo cp cuda/include/cudnn.h /usr/local/cuda-11.3/include/
sudo cp cuda/lib64/libcudnn* /usr/local/cuda-11.3/lib64/
sudo chmod a+r /usr/local/cuda-11.3/include/cudnn.h
sudo chmod a+r /usr/local/cuda-11.3/lib64/libcudnn*

3.安装TensorRT，从官网下载，下载tar版本，解压后有lib64文件夹，路径加入环境变量，进入python文件夹安装whl

注意事项:
1.原始CUDA但版本不对，一定要匹配，要不麻烦死，卸载nvidia，出错可以去“软件和更新”那其他软件那去掉一下相关的对沟，要不卸载不了，显示包不对应

sudo apt-get remove --auto-remove nvidia-cuda-toolkit

2.查看CUDA版本

cat /usr/local/cuda-11.3/version.json

3.添加环境变量

vim ~/.bashrc

加入

export PATH="/usr/local/cuda-11.3/bin:$PATH"
	export LD_LIBRARY_PATH="/usr/local/cuda-11.3/lib64:$LD_LIBRARY_PATH"
	export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/nercar/TensorRT-7.0.0.11/lib
	export CUDA_INSTALL_DIR=/usr/local/cuda-11.3/lib64
	export CUDNN_INSTALL_DIR=/usr/local/cuda-11.3/lib64

执行

source ~/.bashrc
nvcc -V

4.pycharm上出错没有so文件时，通过在pycharm配置环境变量解决，run里边配置环境

PYTHONUNBUFFERED=1;CUDNN_HOME=/usr/local/cuda-11.3/lib64;LD_LIBRARY_PATH=/usr/local/cuda-11.3/lib64\;/home/nercar/TensorRT-8.0.1.6/lib

5.tensorrt8以上会出现有些东西未声明，可以参照自带的历程

实现bisenet的推理

pth转onnx

import argparse
import torch
from lib.models import model_factory
from configs import set_cfg_from_file
import sys
sys.path.insert(0, '.')

torch.set_grad_enabled(False)
parse = argparse.ArgumentParser()
parse.add_argument('--config', dest='config', type=str, default='../configs/bisenetv2_city.py', )
parse.add_argument('--weight-path', dest='weight_pth', type=str, default='model/bisenetv2.pth')
parse.add_argument('--outpath', dest='out_pth', type=str, default='model/bisenetv2.onnx')
parse.add_argument('--aux-mode', dest='aux_mode', type=str, default='pred')
args = parse.parse_args()

cfg = set_cfg_from_file(args.config)
if cfg.use_sync_bn: cfg.use_sync_bn = False

net = model_factory[cfg.model_type](cfg.n_cats, aux_mode=args.aux_mode)
net.load_state_dict(torch.load(args.weight_pth, map_location='cpu'), strict=False)
net.eval()

dummy_input = torch.randn(1, 3, 1024, 2048)
input_names = ['input_image']
output_names = ['preds', ]
torch.onnx.export(net, dummy_input, args.out_pth,
                  input_names=input_names,
                  output_names=output_names,
                  verbose=False,
                  opset_version=11)

onnx转trt

from __future__ import print_function
import os
import sys
import cv2
import time
import common
import numpy as np
import tensorrt as trt
sys.path.insert(1, os.path.join(sys.path[0], ".."))

TRT_LOGGER = trt.Logger()


def get_engine(onnx_file_path, engine_file_path=""):
    def build_engine():
        with trt.Builder(TRT_LOGGER) as builder, \
                builder.create_network(common.EXPLICIT_BATCH) as network, \
                builder.create_builder_config() as config, \
                trt.OnnxParser(network, TRT_LOGGER) as parser, \
                trt.Runtime(TRT_LOGGER) as runtime:
            config.max_workspace_size = 1 << 30  # 256MiB
            builder.max_batch_size = 1
            if not os.path.exists(onnx_file_path):
                print('ONNX file {} not found.'.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
            network.get_input(0).shape = [1, 3, 512, 1024]
            print('Completed parsing of ONNX file')
            print('Building an engine from file {}; this may take a while...'.format(onnx_file_path))
            plan = builder.build_serialized_network(network, config)
            engine = runtime.deserialize_cuda_engine(plan)
            print("Completed creating Engine")
            with open(engine_file_path, "wb") as f:
                f.write(plan)
            return engine

    if os.path.exists(engine_file_path):
        print("Reading engine from file {}".format(engine_file_path))
        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
            return runtime.deserialize_cuda_engine(f.read())
    else:
        return build_engine()


mean = (0.3257, 0.3690, 0.3223)
std = (0.2112, 0.2148, 0.2115)


def main():
    onnx_file_path = 'model/bisenetv2.onnx'
    engine_file_path = "model/bisenetv2.trt"
    input_image_path = 'data/1.png'

    image = cv2.imread(input_image_path, 1)
    image = cv2.resize(image, (1024, 512))
    image = image/255.
    # image = ((image / 255.0) - mean) / std
    image = np.transpose(image, [2, 0, 1])
    image = np.expand_dims(image, axis=0)
    image = np.array(image, dtype=np.float32, order='C')

    with get_engine(onnx_file_path, engine_file_path) as engine, engine.create_execution_context() as context:
        inputs, outputs, bindings, stream = common.allocate_buffers(engine)
        print('Running inference on image {}...'.format(input_image_path))
        inputs[0].host = image
        start = time.time()
        trt_outputs = common.do_inference_v2(context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream)
        end = time.time()
        print("use time:", end-start)
        # stream.synchronize()
    dst = trt_outputs[0].reshape((1024, 2048))
    cv2.namedWindow("dst", 0)
    cv2.imshow("dst", np.array(dst, np.uint8) * 10)
    cv2.waitKey(0)


if __name__ == '__main__':
    main()

common

import argparse
import os
import time
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)


def GiB(val):
    return val * 1 << 30


def add_help(description):
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    args, _ = parser.parse_known_args()


def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""):
    # Standard command-line arguments for all samples.
    kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data")
    parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-d", "--datadir",
                        help="Location of the TensorRT sample data directory, and any additional data directories.",
                        action="append", default=[kDEFAULT_DATA_ROOT])
    args, _ = parser.parse_known_args()

    def get_data_path(data_dir):
        # If the subfolder exists, append it to the path, otherwise use the provided path as-is.
        data_path = os.path.join(data_dir, subfolder)
        if not os.path.exists(data_path):
            if data_dir != kDEFAULT_DATA_ROOT:
                print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.")
            data_path = data_dir
        # Make sure data directory exists.
        if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT:
            print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(
                data_path))
        return data_path

    data_paths = [get_data_path(data_dir) for data_dir in args.datadir]
    return data_paths, locate_files(data_paths, find_files, err_msg)


def locate_files(data_paths, filenames, err_msg=""):
    found_files = [None] * len(filenames)
    for data_path in data_paths:
        # Find all requested files.
        for index, (found, filename) in enumerate(zip(found_files, filenames)):
            if not found:
                file_path = os.path.abspath(os.path.join(data_path, filename))
                if os.path.exists(file_path):
                    found_files[index] = file_path

    # Check that all files were found
    for f, filename in zip(found_files, filenames):
        if not f or not os.path.exists(f):
            raise FileNotFoundError(
                "Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg))
    return found_files


# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem

    def __str__(self):
        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)

    def __repr__(self):
        return self.__str__()


# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        # Allocate host and device buffers
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        # Append the device buffer to device bindings.
        bindings.append(int(device_mem))
        # Append to the appropriate list.
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream


def do_inference(context, bindings, inputs, outputs, stream, batch_size=1):
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    stream.synchronize()
    return [out.host for out in outputs]


def do_inference_v2(context, bindings, inputs, outputs, stream):
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    # stream.synchronize()
    return [out.host for out in outputs]

这篇关于Ubuntu20.04配置TensorRT的文章就介绍到这儿，希望我们推荐的文章对大家有所帮助，也希望大家多多支持为之网！