Introduction

Base stone of AI_power, maintain all inference of AI_Power models.

Wrapper

• Supply different model infer wrapper, including ONNX/TensorRT/Torch JIT;
• Support onnx different Execution Providers (EP) , including cpu/gpu/trt/trt16/int8;
• High level mmlab model (converted) infer wrapper, including MMPose/MMDet;

Model Convert

• torch2jit torch2onnx etc.
• detectron2 to onnx
• modelscope to onnx
• onnx2simple2trt
• tf2pb2onnx

Model Tools

• torch model edit
• onnx model shape/speed test (different EP)
• common scripts from onnxruntime

Usage

onnx model speed test

from apstone import ONNXModel

onnx_p = 'pretrain_models/sr_lib/realesr-general-x4v3-dynamic.onnx'
input_dynamic_shape = (1, 3, 96, 72)  # None
# cpu gpu trt trt16 int8
ONNXModel(onnx_p, provider='cpu', debug=True, input_dynamic_shape=input_dynamic_shape).speed_test()

Install

pip install apstone

Envs

Execution Providers	Needs
cpu	pip install onnxruntime
gpu	pip install onnxruntime-gpu
trt/trt16/int8	onnxruntime-gpu compiled with tensorrt EP
TensorRT	pip install tensorrt pycuda
torch JIT	install pytorch