flambe.nn.rnn

Module Contents

flambe.nn.rnn.logger[source]
class flambe.nn.rnn.RNNEncoder(input_size: int, hidden_size: int, n_layers: int = 1, rnn_type: str = 'lstm', dropout: float = 0, bidirectional: bool = False, layer_norm: bool = False, highway_bias: float = 0, rescale: bool = True, enforce_sorted: bool = False, **kwargs)[source]

Bases: flambe.nn.module.Module

Implements a multi-layer RNN.

This module can be used to create multi-layer RNN models, and provides a way to reduce to output of the RNN to a single hidden state by pooling the encoder states either by taking the maximum, average, or by taking the last hidden state before padding.

Padding is dealt with by using torch’s PackedSequence.

rnn

The rnn submodule

Type:nn.Module
forward(self, data: Tensor, state: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None)[source]

Performs a forward pass through the network.

Parameters:
  • data (Tensor) – The input data, as a float tensor of shape [B x S x E]
  • state (Tensor) – An optional previous state of shape [L x B x H]
  • padding_mask (Tensor, optional) – The padding mask of shape [B x S], dtype should be bool
Returns:

  • Tensor – The encoded output, as a float tensor of shape [B x S x H]
  • Tensor – The encoded state, as a float tensor of shape [L x B x H]

class flambe.nn.rnn.PooledRNNEncoder(input_size: int, hidden_size: int, n_layers: int = 1, rnn_type: str = 'lstm', dropout: float = 0, bidirectional: bool = False, layer_norm: bool = False, highway_bias: float = 0, rescale: bool = True, pooling: str = 'last')[source]

Bases: flambe.nn.module.Module

Implement an RNNEncoder with additional pooling.

This class can be used to obtan a single encoded output for an input sequence. It also ignores the state of the RNN.

forward(self, data: Tensor, state: Optional[Tensor] = None, padding_mask: Optional[Tensor] = None)[source]

Perform a forward pass through the network.

Parameters:
  • data (torch.Tensor) – The input data, as a float tensor of shape [B x S x E]
  • state (Tensor) – An optional previous state of shape [L x B x H]
  • padding_mask (Tensor, optional) – The padding mask of shape [B x S]
Returns:

The encoded output, as a float tensor of shape [B x H]

Return type:

torch.Tensor