MaskedGlobalPoolingBlock¶
-
class
maze.perception.blocks.general.masked_global_pooling.
MaskedGlobalPoolingBlock
(*args: Any, **kwargs: Any)¶ A block applying global pooling with optional masking. Pooling is applied w.r.t. the selected pooling function, the pooling dimension and the optional mask (in_keys[1]).
- Parameters
in_keys – One key identifying the input tensor, and an optional second one identifying the mask tensor.
out_keys – One key identifying the output tensors.
in_shapes – List of input shapes.
pooling_func – Options: {‘mean’, ‘sum’, ‘max’}.
pooling_dim – The dimension(s) along which the pooling functions get applied.
-
forward
(block_input: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor]¶ (overrides
PerceptionBlock
)Forward pass through the block, iterating over the first 2 dimensions and pooling the rest in dim=0.
- param block_input
The block’s input dictionary.
- return
The block’s output dictionary.