128 lines
3.8 KiB
Python
128 lines
3.8 KiB
Python
import tensorrt as trt
|
|
from cuda.bindings import driver as cuda, runtime as cudart
|
|
import numpy as np
|
|
def cuda_call(call):
|
|
err, res = call[0], call[1:]
|
|
check_cuda_err(err)
|
|
if len(res) == 1:
|
|
res = res[0]
|
|
return res
|
|
|
|
|
|
def check_cuda_err(err):
|
|
if isinstance(err, cuda.CUresult):
|
|
if err != cuda.CUresult.CUDA_SUCCESS:
|
|
raise RuntimeError("Cuda Error: {}".format(err))
|
|
if isinstance(err, cudart.cudaError_t):
|
|
if err != cudart.cudaError_t.cudaSuccess:
|
|
raise RuntimeError("Cuda Runtime Error: {}".format(err))
|
|
else:
|
|
raise RuntimeError("Unknown error type: {}".format(err))
|
|
|
|
|
|
|
|
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
|
|
nbytes = host_arr.size * host_arr.itemsize
|
|
cuda_call(
|
|
cudart.cudaMemcpy(
|
|
host_arr, device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost
|
|
)
|
|
)
|
|
|
|
|
|
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
|
|
nbytes = host_arr.size * host_arr.itemsize
|
|
cuda_call(
|
|
cudart.cudaMemcpy(
|
|
device_ptr, host_arr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice
|
|
)
|
|
)
|
|
|
|
|
|
|
|
class TensorRTModel:
|
|
def __init__(self, engine_path):
|
|
logger = trt.Logger(trt.Logger.ERROR)
|
|
trt.init_libnvinfer_plugins(logger, namespace="")
|
|
with open(engine_path, "rb") as f, trt.Runtime(logger) as runtime:
|
|
assert runtime
|
|
engine = runtime.deserialize_cuda_engine(f.read())
|
|
assert engine
|
|
context = engine.create_execution_context()
|
|
assert context
|
|
|
|
self.engine = engine
|
|
self.context = context
|
|
self.logger = logger
|
|
|
|
def prepare(self):
|
|
engine = self.engine
|
|
inputs = []
|
|
outputs = []
|
|
allocations = []
|
|
for i in range(engine.num_io_tensors):
|
|
name = engine.get_tensor_name(i)
|
|
is_input = False
|
|
if engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
|
|
is_input = True
|
|
dtype = engine.get_tensor_dtype(name)
|
|
shape = engine.get_tensor_shape(name)
|
|
if is_input:
|
|
batch_size = shape[0]
|
|
size = np.dtype(trt.nptype(dtype)).itemsize
|
|
for s in shape:
|
|
size *= s
|
|
allocation = cuda_call(cudart.cudaMalloc(size))
|
|
binding = {
|
|
"index": i,
|
|
"name": name,
|
|
"dtype": np.dtype(trt.nptype(dtype)),
|
|
"shape": list(shape),
|
|
"allocation": allocation,
|
|
"size": size,
|
|
}
|
|
allocations.append(allocation)
|
|
if is_input:
|
|
inputs.append(binding)
|
|
else:
|
|
outputs.append(binding)
|
|
|
|
assert batch_size > 0
|
|
assert len(inputs) > 0
|
|
assert len(outputs) > 0
|
|
assert len(allocations) > 0
|
|
|
|
self.batch_size = batch_size
|
|
self.inputs = inputs
|
|
self.outputs = outputs
|
|
self.allocations = allocations
|
|
|
|
assert(len(self.inputs) == 1)
|
|
|
|
def allocate_outputs(self):
|
|
outputs = self.outputs
|
|
cpu_allocs = dict()
|
|
|
|
for output in outputs:
|
|
dtype = output['dtype']
|
|
shape = output['shape']
|
|
name = output['name']
|
|
output.update({'allocated': np.empty(shape, dtype=dtype)})
|
|
|
|
self.cpu_allocs = cpu_allocs
|
|
|
|
def score(self, data_ptr):
|
|
curr_alloc = self.allocations.copy()
|
|
curr_alloc[0] = data_ptr
|
|
self.context.execute_v2(curr_alloc)
|
|
|
|
results = dict()
|
|
for c_out in self.outputs:
|
|
c_name = c_out['name']
|
|
c_alloc_idx = c_out['index']
|
|
c_alloc_array = c_out['allocated']
|
|
memcpy_device_to_host(c_alloc_array, curr_alloc[c_alloc_idx])
|
|
results[c_name] = c_alloc_array
|
|
|
|
return results
|