ONNX export is working

1 files changed, 173 insertions, 0 deletions

diff --git a/run_engine.py b/run_engine.py
new file mode 100644
index 0000000..700494d
--- /dev/null
+++ b/run_engine.py
@@ -0,0 +1,173 @@
+import os
+import sys
+from random import randint
+import numpy as np
+import tensorrt
+
+try:
+    from PIL import Image
+    import pycuda.driver as cuda
+    import pycuda.gpuarray as gpuarray
+    import pycuda.autoinit
+    import argparse
+except ImportError as err:
+    sys.stderr.write("""ERROR: failed to import module ({})
+Please make sure you have pycuda and the example dependencies installed.
+https://wiki.tiker.net/PyCuda/Installation/Linux
+pip(3) install tensorrt[examples]
+""".format(err))
+    exit(1)
+
+try:
+    import tensorrt as trt
+    from tensorrt.parsers import caffeparser
+    from tensorrt.parsers import onnxparser    
+except ImportError as err:
+    sys.stderr.write("""ERROR: failed to import module ({})
+Please make sure you have the TensorRT Library installed
+and accessible in your LD_LIBRARY_PATH
+""".format(err))
+    exit(1)
+
+
+G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.INFO)
+
+class Profiler(trt.infer.Profiler):
+    """
+    Example Implimentation of a Profiler
+    Is identical to the Profiler class in trt.infer so it is possible
+    to just use that instead of implementing this if further
+    functionality is not needed
+    """
+    def __init__(self, timing_iter):
+        trt.infer.Profiler.__init__(self)
+        self.timing_iterations = timing_iter
+        self.profile = []
+
+    def report_layer_time(self, layerName, ms):
+        record = next((r for r in self.profile if r[0] == layerName), (None, None))
+        if record == (None, None):
+            self.profile.append((layerName, ms))
+        else:
+            self.profile[self.profile.index(record)] = (record[0], record[1] + ms)
+
+    def print_layer_times(self):
+        totalTime = 0
+        for i in range(len(self.profile)):
+            print("{:40.40} {:4.3f}ms".format(self.profile[i][0], self.profile[i][1] / self.timing_iterations))
+            totalTime += self.profile[i][1]
+        print("Time over all layers: {:4.2f} ms per iteration".format(totalTime / self.timing_iterations))
+
+
+def get_input_output_names(trt_engine):
+    nbindings = trt_engine.get_nb_bindings();
+    maps = []
+
+    for b in range(0, nbindings):
+        dims = trt_engine.get_binding_dimensions(b).to_DimsCHW()
+        name = trt_engine.get_binding_name(b)
+        type = trt_engine.get_binding_data_type(b)
+        
+        if (trt_engine.binding_is_input(b)):
+            maps.append(name)
+            print("Found input: ", name)
+        else:
+            maps.append(name)
+            print("Found output: ", name)
+
+        print("shape=" + str(dims.C()) + " , " + str(dims.H()) + " , " + str(dims.W()))
+        print("dtype=" + str(type))
+    return maps
+
+def create_memory(engine, name,  buf, mem, batchsize, inp, inp_idx):
+    binding_idx = engine.get_binding_index(name)
+    if binding_idx == -1:
+        raise AttributeError("Not a valid binding")
+    print("Binding: name={}, bindingIndex={}".format(name, str(binding_idx)))
+    dims = engine.get_binding_dimensions(binding_idx).to_DimsCHW()
+    eltCount = dims.C() * dims.H() * dims.W() * batchsize
+
+    if engine.binding_is_input(binding_idx):
+        h_mem = inp[inp_idx]
+        inp_idx = inp_idx + 1
+    else:
+        h_mem = np.random.uniform(0.0, 255.0, eltCount).astype(np.dtype('f4'))
+
+    d_mem = cuda.mem_alloc(eltCount * 4)
+    cuda.memcpy_htod(d_mem, h_mem)
+    buf.insert(binding_idx, int(d_mem))
+    mem.append(d_mem)
+    return inp_idx
+
+
+#Run inference on device
+def time_inference(engine, batch_size, inp):
+    bindings = []
+    mem = []
+    inp_idx = 0
+    for io in get_input_output_names(engine):
+        inp_idx = create_memory(engine, io,  bindings, mem,
+                                batch_size, inp, inp_idx)
+
+    context = engine.create_execution_context()
+    g_prof = Profiler(500)
+    context.set_profiler(g_prof)
+    for i in range(iter):
+        context.execute(batch_size, bindings)
+    g_prof.print_layer_times()
+    
+    context.destroy() 
+    return
+
+
+def convert_to_datatype(v):
+    if v==8:
+        return trt.infer.DataType.INT8
+    elif v==16:
+        return trt.infer.DataType.HALF
+    elif v==32:
+        return trt.infer.DataType.FLOAT
+    else:
+        print("ERROR: Invalid model data type bit depth: " + str(v))
+        return trt.infer.DataType.INT8
+
+def run_trt_engine(engine_file, bs, it):
+    engine = trt.utils.load_engine(G_LOGGER, engine_file)
+    time_inference(engine, bs, it)
+
+def run_onnx(onnx_file, data_type, bs, inp):
+    # Create onnx_config
+    apex = onnxparser.create_onnxconfig()
+    apex.set_model_file_name(onnx_file)
+    apex.set_model_dtype(convert_to_datatype(data_type))
+
+     # create parser
+    trt_parser = onnxparser.create_onnxparser(apex)
+    assert(trt_parser)
+    data_type = apex.get_model_dtype()
+    onnx_filename = apex.get_model_file_name()
+    trt_parser.parse(onnx_filename, data_type)
+    trt_parser.report_parsing_info()
+    trt_parser.convert_to_trtnetwork()
+    trt_network = trt_parser.get_trtnetwork()
+    assert(trt_network)
+
+    # create infer builder
+    trt_builder = trt.infer.create_infer_builder(G_LOGGER)
+    trt_builder.set_max_batch_size(max_batch_size)
+    trt_builder.set_max_workspace_size(max_workspace_size)
+    
+    if (apex.get_model_dtype() == trt.infer.DataType_kHALF):
+        print("-------------------  Running FP16 -----------------------------")
+        trt_builder.set_half2_mode(True)
+    elif (apex.get_model_dtype() == trt.infer.DataType_kINT8): 
+        print("-------------------  Running INT8 -----------------------------")
+        trt_builder.set_int8_mode(True)
+    else:
+        print("-------------------  Running FP32 -----------------------------")
+        
+    print("----- Builder is Done -----")
+    print("----- Creating Engine -----")
+    trt_engine = trt_builder.build_cuda_engine(trt_network)
+    print("----- Engine is built -----")
+    time_inference(engine, bs, inp)