MaskedGlobalPoolingBlock

class maze.perception.blocks.general.masked_global_pooling.MaskedGlobalPoolingBlock(*args: Any, **kwargs: Any)

A block applying masked global pooling. Pooling is applied wrt the mask (in_keys[1]) and the selected pooling function. That is, in the forward pass the input tensor 1 is iterated over in the first 2 dimensions, where the elements are selected based on the mask, before applying the pooling function.

Parameters

• in_keys – One key identifying the input tensors.

• out_keys – One key identifying the output tensors.

• in_shapes – List of input shapes.

• pooling_func – Options: {‘mean’}. So far only mean pooling is supported.

• pooling_dim – The dimension(s) along which the pooling functions get applied.

forward(block_input: Dict[str, torch.Tensor]) → Dict[str, torch.Tensor]

Forward pass through the block, iterating over the first 2 dimensions and pooling the rest in dim=0.

Parameters

block_input – The block’s input dictionary.

Returns

The block’s output dictionary.