collect_activation_statistics

collect_activation_statistics(model: nn.Module, observations: np.ndarray | torch.Tensor, *, device: torch.device, batch_size: int = 256, near_dead_threshold: float = 0.01, max_rank_rows: int = 4096) -> list[LayerActivationStatistics]

Run observations through a model and measure every activation layer.

Reports, per layer, how many units are dead or near-dead and the effective rank of what the layer emits. Together those say how much of the network's nominal width is doing anything. A long RL run usually loses capacity this way rather than by being too small to begin with, so compare the trained model against a freshly initialised one of the same shape before reading a number as high or low.

Parameters:

• model (Module) – The model. Left in whatever train/eval mode it arrived in.
• observations (ndarray | Tensor) – Observations in the model's input dtype and shape, (N, *observation_shape); a numpy array or a tensor (on any device).
• device (device) – Device to run the forward passes on.
• batch_size (int) – Observations per forward pass.
• near_dead_threshold (float) – A rectifier unit active on less than this fraction of inputs counts as near-dead.
• max_rank_rows (int) – Rows of activations kept per layer for the rank estimate, spread over the batches.

Returns:

• list[LayerActivationStatistics] – One entry per activation layer, in forward order.

Source