ValidationManager#

The definition of ValidationManager class.

class orbax.export.validate.validation_manager.ValidationManager(module, serving_configs, model_inputs)[source][source]#

Validate the JaxModule and its output tf saved model.

__init__(module, serving_configs, model_inputs)[source][source]#

Create the ValidationManager ojbect.

Parameters:

module (Union[JaxModule, Mapping[str, Callable[[Any], Any]]]) – the JaxModule object.
serving_configs (Sequence[ServingConfig]) – the ServingConfig Sequence.
model_inputs (Union[Sequence[Any], Mapping[str, Sequence[Any]]]) – The inputs for saved TF SavedModel. It support two formats: (1) A mapping of signature key to a sequences batch inputs; or (2) a sequence of batch inputs to validate all signatures.

validate(loaded_model, with_xprof=False, report_option=None)[source][source]#

Validates the baseline and candidate function map.

classmethod check_input(inputs, batch_mode=True)[source][source]#

check model input format.

Parameters:

inputs (Union[Any, Sequence[Any]]) – model inputs. If batch_mode == True, inputs should be a list.
batch_mode (bool) – it decide inputs is a list of the input or a single input.

Return type:

None

classmethod check_output(baseline_result, candidate_result)[source][source]#

check model output format.

ValidationManager