import onnxruntime import torch providers = [ # ('TensorrtExecutionProvider', { # 'device_id': 0, # 'trt_max_workspace_size': 8 * 1024 * 1024 * 1024, # 'trt_fp16_enable': True, # 'trt_engine_cache_enable': True, # }), ('CUDAExecutionProvider', { 'device_id': 0, 'arena_extend_strategy': 'kSameAsRequested', 'gpu_mem_limit': 8 * 1024 * 1024 * 1024, 'cudnn_conv_algo_search': 'HEURISTIC', }) ] def load_onnx(file_path: str): assert file_path.endswith(".onnx") sess_opt = onnxruntime.SessionOptions() ort_session = onnxruntime.InferenceSession(file_path, sess_opt=sess_opt, providers=providers) return ort_session def load_onnx_caller(file_path: str, single_output=False): ort_session = load_onnx(file_path) def caller(*args): torch_input = isinstance(args[0], torch.Tensor) if torch_input: torch_input_dtype = args[0].dtype torch_input_device = args[0].device # check all are torch.Tensor and have same dtype and device assert all([isinstance(arg, torch.Tensor) for arg in args]), "All inputs should be torch.Tensor, if first input is torch.Tensor" assert all([arg.dtype == torch_input_dtype for arg in args]), "All inputs should have same dtype, if first input is torch.Tensor" assert all([arg.device == torch_input_device for arg in args]), "All inputs should have same device, if first input is torch.Tensor" args = [arg.cpu().float().numpy() for arg in args] ort_inputs = {ort_session.get_inputs()[idx].name: args[idx] for idx in range(len(args))} ort_outs = ort_session.run(None, ort_inputs) if torch_input: ort_outs = [torch.tensor(ort_out, dtype=torch_input_dtype, device=torch_input_device) for ort_out in ort_outs] if single_output: return ort_outs[0] return ort_outs return caller