aihwkit.utils.analog_info module

Analog Information utility.

This module prints relevant information about the model and its analog execution.

class aihwkit.utils.analog_info.AnalogInfo(model, input_size=None, rpu_config=None)

Bases: object

Class for computing and storing results of the analog summary.

Parameters:

• model (Module) –

• input_size (Any) –

• rpu_config (RPUConfigBase | None) –

calculate_num_analog()

Calculate the number of analog layers.

Return type:

int

calculate_num_tiles()

Calculate the total number of tiles needed by the model.

Return type:

int

create_layer_summary()

Create the layer summary list.

This list contains LayerInfo elements that corresponds to each layer of the model.

Return type:

List[LayerInfo]

register_hooks_recursively(module, hook)

Hooks the function into all layers with no children (or only analog tiles as childrens).

Parameters:

• module (Module) –

• hook (Any) –

Return type:

None

class aihwkit.utils.analog_info.LayerInfo(module, rpu_config=None, input_size=None, output_size=None)

Bases: object

Class for storing layer statistics and information.

Parameters:

• module (Module) –

• rpu_config (RPUConfigBase | None) –

• input_size (Any) –

• output_size (Any) –

calculate_reuse_factor()

Compute the reuse factor.

The reuse factor is the number of vector matrix multiplication a layer computes.

Return type:

None

input_size: Any

isanalog: bool

kernel_size: Any

layer_summary_dict()

Return a dictionary with all layer’s information.

Return type:

dict

module: Module

name: str

num_tiles: int

output_size: Any

reuse_factor: int

set_kernel_size()

Set kernel size attribute.

Return type:

None

set_tiles_info()

Create TileInfo objects for each tile of the layer.

Return type:

List[TileInfo]

tiles_info: List[TileInfo]

tiles_summary_dict()

Return a dictionary with all tiles information.

Return type:

dict

class aihwkit.utils.analog_info.TileInfo(tile, is_mapped)

Bases: object

Class for storing tile statistics and information.

Parameters:

• tile (TileModule) –

• is_mapped (bool) –

log_in_size: Any

log_out_size: Any

phy_in_size: Any

phy_out_size: Any

tile_summary_dict()

Return a dictionary with the tile info.

Return type:

dict

utilization: float

aihwkit.utils.analog_info.analog_summary(model, input_size=None, rpu_config=None)

Summarize the given PyTorch model.

Summarized information includes:

1. Layer names,

2. input/output shapes,

3. kernel shape,

4. # of digital parameters,

5. # of analog parameters,

6. # of analog tiles

7. reuse factor

Parameters:

• model (Module) – PyTorch model to run on the analog platform.

• input_size (Any | None) – required to run a forward pass of the model.

• rpu_config (RPUConfigBase | None) – resistive processing unit configuration.

Returns:

AnalogInfo Object.

Return type:

AnalogInfo