Skip to content

model_tester

model_tester(config, input_dict, output_shape_dict, output_dict_type, output_test_strategy=None, output_equivalence_dict=None)

This functions tests the mdoel using different checks. In order:

  • check that model's outputs shapes matched output_shape_dict
  • check that the model can be trained, we use output_dict_type to know how to run output's specific tests. For example, for image classification you should pass output_dict_type=ModelForImageClassificationOutput and we will try to run outputs["logits"].mean().backward()
  • optionally, if output_equivalence_dict we will check that model's outputs match

Parameters:

Name Type Description Default
config Config

Configuration we will use to build the model.

required
input_dict Dict[str, Tensor]

Dictionary containing the inputs for the model. E.g. { "pixel_values" : torch.randn((1, 3, 224, 224))}

required
output_shape_dict Dict[str, Tuple[int]]

Dictionary containing the expected shaped in the output. E.g. {"logits": (1, 1000)}

required
output_dict_type TypedDict

The type of output, e.g. ModelForImageClassificationOutput](). Used to now which test strategy to run, based on the type we know how to test it. Defaults to None.

required
output_test_strategy Optional[Callable[[TypedDict]]]

If passed, we will use this strategy instead.

None
output_equivalence_dict Optional[Dict[str, Tensor]]

If passes, we will check that the model's output are equal to the values inside it. Defaults to None.

None
Source code in glasses/utils/model_tester.py
def model_tester(
    config: Config,
    input_dict: Dict[str, Tensor],
    output_shape_dict: Dict[str, Tuple[int]],
    output_dict_type: TypedDict,
    output_test_strategy: Optional[Callable[[TypedDict], None]] = None,
    output_equivalence_dict: Optional[Dict[str, Tensor]] = None,
):
    """
    This functions tests the mdoel using different checks. In order:

    - check that model's outputs shapes matched `output_shape_dict`
    - check that the model can be trained, we use `output_dict_type` to know how to run output's specific tests. For example, for image classification you should pass `output_dict_type=ModelForImageClassificationOutput` and we will try to run `outputs["logits"].mean().backward()`
    - optionally, if `output_equivalence_dict` we will check that model's outputs match

    Args:
        config (Config): Configuration we will use to build the model.
        input_dict (Dict[str, Tensor]): Dictionary containing the inputs for the model. E.g. `{ "pixel_values" : torch.randn((1, 3, 224, 224))}`
        output_shape_dict (Dict[str, Tuple[int]]): Dictionary containing the expected shaped in the output. E.g. ` {"logits": (1, 1000)}`
        output_dict_type (TypedDict): The type of output, e.g. ModelForImageClassificationOutput](). Used to now which test strategy to run, based on the type we know how to test it. Defaults to None.
        output_test_strategy (Optional[Callable[[TypedDict]]], optional): If passed, we will use this strategy instead.
        output_equivalence_dict (Optional[Dict[str, Tensor]], optional): If passes, we will check that the model's output are equal to the values inside it. Defaults to None.
    """
    with torch.no_grad():
        # we are able to create the model from a config
        model = config.build().eval()
        model_output_dict = model(**input_dict)
        check_model_output_dict_shape(model_output_dict, output_shape_dict)
    model = model.train()
    model_output_dict = model(**input_dict)
    # we are able to check that the model can be trained
    output_test_strategy = output_test_strategy
    if output_test_strategy is None:
        output_test_strategy = model_output_test_strategies[output_dict_type]
    output_test_strategy(model_output_dict)
    check_grad_exists(model)
    # we are able to check outputs equivalence if passed
    if output_equivalence_dict:
        with torch.no_grad():
            # here run the model with a known input
            check_output_equivalence(model_output_dict, output_equivalence_dict)