yamle.utils.tracing_utils module

yamle.utils.tracing_utils.forward_shape_hook(module, input, output)

This function is used to cache the input and output shapes of a module.

The shapes will be stored in the module as MODULE_INPUT_SHAPE_KEY and MODULE_OUTPUT_SHAPE_KEY.

Parameters:

• module (torch.nn.Module) – The module to cache the input and output shapes of.

• input (torch.Tensor) – The input to the module.

• output (torch.Tensor) – The output of the module.

Return type:

None

yamle.utils.tracing_utils.get_sample_input_and_target(method, batch_size=None)

This method is used to get the sample input and target of the model.

Return type:

Tuple[Tensor, Tensor]

yamle.utils.tracing_utils.get_input_shape_from_model(model)

This method is used to get the input shape of the model.

Return type:

Tuple[int, ...]

yamle.utils.tracing_utils.get_output_shape_from_model(model)

This method is used to get the output shape of the model.

Return type:

Tuple[int, ...]

yamle.utils.tracing_utils.trace_input_output_shapes(method)

This method is used to trace the input and output shapes of the model.

Additionally, it will name all the modules in the model.

Return type:

None

yamle.utils.tracing_utils.name_all_modules(model)

This method is used to name all the modules in the model.

Return type:

None