FlattenDenseBlock

class maze.perception.blocks.joint_blocks.flatten_dense.FlattenDenseBlock(*args: Any, **kwargs: Any)

A block containing a flattening stage followed by a dense layer block.

For details on flattening see FlattenBlock. For details on dense layers see DenseBlock.

Parameters:

  • in_keys – One key identifying the input tensors.

  • out_keys – One key identifying the output tensors.

  • in_shapes – List of input shapes.

  • num_flatten_dims – the number of dimensions to flatten out (from right).

  • hidden_units – List containing the number of hidden units for hidden layers.

  • non_lin – The non-linearity to apply after each layer.

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

(overrides PerceptionBlock)

implementation of ShapeNormalizationBlock interface