Profiler module

calculate_features_networks.calculate_flops(model, input_tensor)

calculate_features_networks.calculate_params(model)

calculate_features_networks.estimate_memory_inference(model, sample_input, batch_size=1, use_amp=False, device=0)

Predict the maximum memory usage of the model. :param optimizer_type: the class name of the optimizer to instantiate :type optimizer_type: Type :param model: the neural network model :type model: nn.Module :param sample_input: A sample input to the network. It should be

a single item, not a batch, and it will be replicated batch_size times.

Parameters:

• batch_size (int) – the batch size

• use_amp (bool) – whether to estimate based on using mixed precision

• device (torch.device) – the device to use

calculate_features_networks.get_parser(model)

Function to get the parser with the arguments.

Raises:

ValueError – The specified configuration doesn’t exist

Returns:

Arguments from the command line.

Return type:

argparse.Namespace

calculate_features_networks.human_readable(num)

calculate_features_networks.main(name_model, data)

calculate_features_networks.main_profile(name_model, data)

calculate_features_networks.memory_profile(prof, sort_by, row_limit)

calculate_features_networks.parse_memory_size(size_str)

Convert memory size string to bytes.

calculate_features_networks.separate_number_and_unit(s)

calculate_features_networks.time_profile(prof, sort_by, row_limit)

calculate_features_networks.write_data_json(data, net, num_param, num_flops, memory, feature_size, num_heads)