onnx¶

onnx model management

Classes¶

Functions¶

get_prediction_function(args)

Module Contents¶

class onnx.OnnxRunner(session, setup, device)¶

Parameters:

session (onnxruntime.InferenceSession)
setup (rex_xai.input.input_data.Setup)

session¶

input_shape¶

output_shape¶

input_name¶

output_name¶

setup: rex_xai.input.input_data.Setup¶

device: str¶

run_on_cpu(tensors, target, raw, binary_threshold=None)¶

Convert a pytorch tensor, or list of tensors, to numpy arrays on the cpu for onnx inference.

Parameters:

tensors (Union[torch.Tensor, List[torch.Tensor]])
target (Optional[rex_xai.responsibility.prediction.Prediction])
raw (bool)
binary_threshold (Optional[float])

run_with_data_on_device(tensors, device, tsize, binary_threshold, raw=False, device_id=0, target=None)¶

gen_prediction_function()¶

onnx.get_prediction_function(args)¶