Introduction To Importing Caffe, TensorFlow And ONNX Models Into TensorRT Using Python

最新推荐文章于 2022-03-30 16:34:29 发布

谢欣燕

最新推荐文章于 2022-03-30 16:34:29 发布

阅读量104

点赞数

分类专栏：笔记文章标签： caffe tensorflow python

本文链接：https://blog.csdn.net/weixin_43941538/article/details/120853002

版权

笔记专栏收录该内容

70 篇文章 6 订阅

订阅专栏

1 caffe_parser

import random
from PIL import Image
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import sys, os
sys.path.insert(1,os.path.join(sys.path[0],".."))
import common
class ModelData(object):
    MODEL_PATH = "ResNet50_fp32.caffemodel"
    DEPLOY_PATH = "ResNet50_N2.prototxt"
    INPUT_SHAPE = (3,224,224)
    OUTPUT_NAME = "prob"
    DTYPE = trt.float32
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def allocate_buffers(engine):
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(ModelData.DTYPE))
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(ModelData.DTYPE))
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    return h_input, d_input, h_output, d_output, stream
def do_inference(context, h_input, d_input, h_output, d_output, stream):
    cuda.memcpy_htod_async(d_input,h_input,stream)
    context.execute_async(bindings=[int(d_input),int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
def build_engine_caffe(model_file, deploy_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
        builder.max_workspace_size = common.GiB(1)
        model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype = ModelData.DTYPE)
        network.mark_output(model_tensors.find(ModelData.OUTPUT_NAME))
        return builder.build_cuda_engine(network)
def load_normalized_test_case(test_image, pagelocked_buffer):
    def normalize_image(image):
        c, h, w = ModelData.INPUT_SHAPE
        return np.asarray(image.resize((w,h), Image.ANTIALIAS)).transpose([2,0,1]).astype(trt.nptype(ModelData.DTYPE)).ravel()
    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
    return test_image
def main():
    _, data_files = common.find_sample_data(description="Runs a ResNet50 network with a TensorRT inference engine.",
                                            subfolder="resnet50",
                                            find_files=["binoculars.jpeg", "reflex_camera.jpeg", "tabby_tiger_cat.jpg",
                                                        ModelData.MODEL_PATH, ModelData.DEPLOY_PATH,
                                                        "class_labels.txt"])
    test_images = data_files[0:3]
    caffe_model_file, caffe_deploy_file, labels_file = data_files[3:]
    labels = open(labels_file, 'r').read().split('\n')
    with build_engine_caffe(caffe_model_file, caffe_deploy_file) as engine:
        h_input, d_input, h_output, d_output, stream = allocate_buffers(engine)
        with engine.create_execution_context() as context:
            test_image = random.choice(test_images)
            test_case = load_normalized_test_case(test_image, h_input)
            do_inference(context, h_input, d_input, h_output, d_output, stream)
            pred = labels[np.argmax(h_output)]
            if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
                print("Correctly recognized " + test_case + " as " + pred)
            else:
                print("Incorrectly recognized " + test_case + " as " + pred)
if __name__ == '__main__':
    main()

2 onnx_parser

import random
from PIL import Image
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import sys, os
sys.path.insert(1,os.path.join(sys.path[0],".."))
import common
class ModelData(object):
    MODEL_PATH = "ResNet50_fp32.caffemodel"
    DEPLOY_PATH = "ResNet50_N2.prototxt"
    INPUT_SHAPE = (3,224,224)
    OUTPUT_NAME = "prob"
    DTYPE = trt.float32
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def allocate_buffers(engine):
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(ModelData.DTYPE))
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(ModelData.DTYPE))
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    return h_input, d_input, h_output, d_output, stream
def do_inference(context, h_input, d_input, h_output, d_output, stream):
    cuda.memcpy_htod_async(d_input,h_input,stream)
    context.execute_async(bindings=[int(d_input),int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
def build_engine_caffe(model_file, deploy_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.CaffeParser() as parser:
        builder.max_workspace_size = common.GiB(1)
        model_tensors = parser.parse(deploy=deploy_file, model=model_file, network=network, dtype = ModelData.DTYPE)
        network.mark_output(model_tensors.find(ModelData.OUTPUT_NAME))
        return builder.build_cuda_engine(network)
def load_normalized_test_case(test_image, pagelocked_buffer):
    def normalize_image(image):
        c, h, w = ModelData.INPUT_SHAPE
        return np.asarray(image.resize((w,h), Image.ANTIALIAS)).transpose([2,0,1]).astype(trt.nptype(ModelData.DTYPE)).ravel()
    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
    return test_image
def main():
    _, data_files = common.find_sample_data(description="Runs a ResNet50 network with a TensorRT inference engine.",
                                            subfolder="resnet50",
                                            find_files=["binoculars.jpeg", "reflex_camera.jpeg", "tabby_tiger_cat.jpg",
                                                        ModelData.MODEL_PATH, ModelData.DEPLOY_PATH,
                                                        "class_labels.txt"])
    test_images = data_files[0:3]
    caffe_model_file, caffe_deploy_file, labels_file = data_files[3:]
    labels = open(labels_file, 'r').read().split('\n')
    with build_engine_caffe(caffe_model_file, caffe_deploy_file) as engine:
        h_input, d_input, h_output, d_output, stream = allocate_buffers(engine)
        with engine.create_execution_context() as context:
            test_image = random.choice(test_images)
            test_case = load_normalized_test_case(test_image, h_input)
            do_inference(context, h_input, d_input, h_output, d_output, stream)
            pred = labels[np.argmax(h_output)]
            if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
                print("Correctly recognized " + test_case + " as " + pred)
            else:
                print("Incorrectly recognized " + test_case + " as " + pred)
if __name__ == '__main__':
    main()

3 uff_parser

import random
from PIL import  Image
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
import sys,os
sys.path.insert(1, os.path.join(sys.path[0],".."))
import common
class ModelData(object):
    MODEL_PATH = 'resnet50-infer-5.uff'
    INPUT_NAME = 'input'
    INPUT_SHAPE = (3,224,224)
    OUTPUT_NAME = "output"
    DTYPE = trt.float32
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
def allocate_buffers(engine):
    h_input = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(0)), dtype=trt.nptype(ModelData.DTYPE))
    h_output = cuda.pagelocked_empty(trt.volume(engine.get_binding_shape(1)), dtype=trt.nptype(ModelData.DTYPE))
    d_input = cuda.mem_alloc(h_input.nbytes)
    d_output = cuda.mem_alloc(h_output.nbytes)
    stream = cuda.Stream()
    return h_input, d_input, h_output, d_output, stream
def do_inference(context, h_input, d_input, h_output, d_output, stream):
    cuda.memcpy_htod_async(d_input, h_output, stream)
    context.execute_async(bindings = [int(d_input), int(d_output)], stream_handle=stream.handle)
    cuda.memcpy_dtoh_async(h_output, d_output, stream)
    stream.synchronize()
def build_engine_uff(model_file):
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network, trt.UffParser() as parser:
        builder.max_workspace_size = common.GiB(1)
        parser.register_input(ModelData.INPUT_NAME, ModelData.INPUT_SHAPE)
        parser.register_output(ModelData.OUTPUT_NAME)
        parser.parse(model_file, network)
        return builder.build_cuda_engine(network)
def load_normalized_test_case(test_image, pagelocked_buffer):
    def normalize_image(image):
        c, h, w = ModelData.INPUT_SHAPE
        return np.asarray(image.resize((w, h), Image.ANTIALIAS)).transpose([2, 0, 1]).astype(
            trt.nptype(ModelData.DTYPE)).ravel()
    np.copyto(pagelocked_buffer, normalize_image(Image.open(test_image)))
    return test_image
def main():
    _, data_files = common.find_sample_data(description="Runs a ResNet50 network with a TensorRT inference engine.",
                                            subfolder="resnet50",
                                            find_files=["binoculars.jpeg", "reflex_camera.jpeg", "tabby_tiger_cat.jpg",
                                                        ModelData.MODEL_PATH, "class_labels.txt"])
    test_images = data_files[0:3]
    uff_model_file, labels_file = data_files[3:]
    labels = open(labels_file, 'r').read().split('\n')
    with build_engine_uff(uff_model_file) as engine:
        h_input, d_input, h_output, d_output, stream = allocate_buffers(engine)
        with engine.create_execution_context() as context:
            test_image = random.choice(test_images)
            test_case = load_normalized_test_case(test_image, h_input)
            do_inference(context, h_input, d_input, h_output, d_output, stream)
            pred = labels[np.argmax(h_output)]
            if "_".join(pred.split()) in os.path.splitext(os.path.basename(test_case))[0]:
                print("Correctly recognized " + test_case + " as " + pred)
            else:
                print("Incorrectly recognized " + test_case + " as " + pred)
if __name__ == '__main__':
    main()

谢欣燕

关注

0
点赞
踩
0

收藏

觉得还不错? 一键收藏
0
评论
Introduction To Importing Caffe, TensorFlow And ONNX Models Into TensorRT Using Python

1 caffe_parserimport randomfrom PIL import Imageimport numpy as npimport pycuda.driver as cudaimport pycuda.autoinitimport tensorrt as trtimport sys, ossys.path.insert(1,os.path.join(sys.path[0],".."))import commonclass ModelData(object): MOD
复制链接

扫一扫