ONNX inference

class torchok.tasks.onnx.ONNXTask(hparams: DictConfig, path_to_onnx: str, providers, keys_mapping_onnx2dataset: Dict[str, str], **kwargs)

A class for ONNX task. This task works only in inference mode.

str_type2numpy_type = {'tensor(double)': 'float64', 'tensor(float)': 'float32', 'tensor(float16)': 'float16', 'tensor(int16)': 'int16', 'tensor(int32)': 'int32', 'tensor(int64)': 'int64', 'tensor(int8)': 'int8', 'tensor(uint8)': 'uint8'}

__init__(hparams: DictConfig, path_to_onnx: str, providers, keys_mapping_onnx2dataset: Dict[str, str], **kwargs)

Init ONNXTask.

Parameters

hparams – Hyperparameters that set in yaml file.
path_to_onnx – path to ONNX model file.
providers – Optional sequence of providers in order of decreasing precedence. Values can either be provider names or tuples of (provider name, options dict). If not provided, then all available providers are used with the default precedence.
keys_mapping_onnx2dataset – mapping from input name of the ONNX model to the input name in the Dataset.

forward_with_gt(batch: Dict[str, Any]) → Dict[str, Tensor]: Is not supported

foward_infer(batch: Dict[str, Tensor]) → Dict[str, Tensor]

Forward onnx model.

Parameters: batch – Dictionary with the items specified in Dataset
Returns: Dictionary with the output items that are specified in model onnx file.

forward_infer_with_gt(batch: Dict[str, Any]) → Dict[str, Tensor]: Forward method for test stage.

test_step(batch: Dict[str, Union[Tensor, int]], batch_idx: int) → None: Complete test loop.

predict_step(batch: Dict[str, Any], batch_idx: int) → Dict[str, Tensor]: Complete predict loop.