onnx_diagnostic.tasks

All submodules contains the three following functions:

  • reduce_model_config(config) -> kwargs: updates the configuration to get a smaller model more suitable for unit tests

  • random_input_kwargs(config) -> kwargs, get_inputs: produces values get_inputs can take to generate dummy inputs suitable for a model defined by its configuration

  • get_inputs(model, config, *args, add_second_input=False, **kwargs) -> dict(inputs=..., dynamic_shapes=...): generates the dummy inputs and dynamic shapes for a specific model and configuration, if add_second_input is True, the function should return a different set of inputs, with different values for the dynamic dimension. This is usually better to rely on the function as the dynamic dimensions may be correlated.

For a specific task, you would write:

kwargs, get_inputs = random_input_kwargs(config)
dummies = get_inputs(model, config, **kwargs)

Or:

from onnx_diagnostic.tasks import random_input_kwargs

kwargs, get_inputs = random_input_kwargs(config, task)  # "text-generation" for example
dummies = get_inputs(model, config, **kwargs)

modules

onnx_diagnostic.tasks.random_input_kwargs(config: Any, task: str) → Tuple[Dict[str, Any], Callable][source][source]

Inputs kwargs. If the configuration is None, the function selects typical dimensions. It returns parameters and a function. The function creates dummy inputs if it receives the parameters returned as a first result.

onnx_diagnostic.tasks.reduce_model_config(config: Any, task: str) → Dict[str, Any][source][source]

Reduces a model size.

onnx_diagnostic.tasks.supported_tasks() → List[str][source][source]

Returns the list of supported tasks.