output_sizes

hybrid_learning.concepts.models.model_extension.output_sizes(model, input_size, layer_ids=None, has_batch_dim=False, ignore_non_tensor_outs=True)[source]

Obtain the output sizes of the given or all layers for given input size. The layer outputs are obtained by feeding the model an all zero Tensor of given input size. The output size can only be determined, if the output of the layer is a torch.Tensor. Other entries are skipped, if ignore_non_tensor_outs=True, and otherwise, an exception is raised. In such cases, consider using dummy_output() directly.

Parameters

  • ignore_non_tensor_outs (bool) –

  • input_size (Sequence[int]) – input size of one sample to feed in; it is assumed to have no batch dimension, if has_batch_dim == False

  • model (Module) – the model to investigate

  • layer_ids (Optional[Sequence[str]]) – the layers to investigate; defaults to all listed in the model’s named_modules

  • has_batch_dim (bool) – whether the given tensor has batch dimension; if not, it is added

Returns

a dict of the sizes for each layer output tensor (batch dimension stripped)

Raises

AttributeError, if one of the considered layers does not output a tensor (but e.g. a dict) and ignore_non_tensor_outs == False; use dummy_output() directly in this case

Return type

Dict[str, Size]