estimate_model_footprint

behavenet.models.ae_model_architecture_generator.estimate_model_footprint(model, input_dim, cutoff_size=20)

Estimate model size to determine if it will fit on a single GPU.

Adapted from http://jacobkimmel.github.io/pytorch_estimating_model_size/.

The estimation: - assumes the model (autoencoder) is symmetric - assumes all values (data/parameters) are float32 - accounts for size of input data - accounts for storage of intermediate layer values and gradients - adds an additional 20% fudge factor

Parameters:

• model (pytorch model) –

• input_dim (array-like) – dimensions of batch with shape (time, n_channels, y_pix, x_pix)

• cutoff_size (float, optional) – terminate estimation once model size grows beyond this point (GB)

Returns:

estimated size of model in bytes

Return type:

int