def test_onnx():
trt = pytrt.Trt()
onnxModel = "../models/model.onnx"
engineFile = ""
customOutput = []
maxBatchSize = 1
calibratorData = [np.ones(28*28)]
mode = 2
trt.CreateEngine( onnxModel, engineFile,customOutput,maxBatchSize,mode,calibratorData)
input_numpy_array = np.zeros(28*28)
trt.DoInference(input_numpy_array) # slightly different from c++
output_numpy_array = trt.GetOutput("Plus214_Output_0")
def export_trt_model(onnxModel, engineFile, input_numpy_array):
r"""
Export a model to trt format.
"""
trt = pytrt.Trt()
customOutput = []
maxBatchSize = 1
calibratorData = []
mode = 2
trt.CreateEngine(onnxModel, engineFile, customOutput, maxBatchSize, mode,
calibratorData)
trt.DoInference(input_numpy_array) # slightly different from c++
return 0
# Apply pre-processing to image.
img = cv2.resize(original_image, (image_width, image_height), interpolation=cv2.INTER_CUBIC)
img = img.astype("float32").transpose(2, 0, 1)[np.newaxis] # (1, 3, h, w)
return img
def normalize(nparray, order=2, axis=-1):
"""Normalize a N-D numpy array along the specified axis."""
norm = np.linalg.norm(nparray, ord=order, axis=axis, keepdims=True)
return nparray / (norm + np.finfo(np.float32).eps)
if __name__ == "__main__":
args = get_parser().parse_args()
trt = pytrt.Trt()
onnxModel = ""
engineFile = args.model_path
customOutput = []
maxBatchSize = 1
calibratorData = []
mode = 2
trt.CreateEngine(onnxModel, engineFile, customOutput, maxBatchSize, mode, calibratorData)
if not os.path.exists(args.output): os.makedirs(args.output)
if args.input:
if os.path.isdir(args.input[0]):
args.input = glob.glob(os.path.expanduser(args.input[0]))
assert args.input, "The input path(s) was not found"