flambe.nn

Package Contents

class flambe.nn.Module[source]

Bases: flambe.compile.Component, torch.nn.Module

Base Flambé Module inteface.

Provides the exact same interface as Pytorch’s nn.Module, but extends it with a useful set of methods to access and clip parameters, as well as gradients.

This abstraction allows users to convert their modules with a single line change, by importing from Flambé instead. Just like every Pytorch module, a forward method should be implemented.

named_trainable_params :Iterator[Tuple[str, nn.Parameter]]

Get all the named parameters with requires_grad=True.

Returns:Iterator over the parameters and their name.
Return type:Iterator[Tuple[str, nn.Parameter]]
trainable_params :Iterator[nn.Parameter]

Get all the parameters with requires_grad=True.

Returns:Iterator over the parameters
Return type:Iterator[nn.Parameter]
gradient_norm :float

Compute the average gradient norm.

Returns:The current average gradient norm
Return type:float
parameter_norm :float

Compute the average parameter norm.

Returns:The current average parameter norm
Return type:float
num_parameters(self, trainable=False)

Gets the number of parameters in the model.

Returns:number of model params
Return type:int
clip_params(self, threshold: float)

Clip the parameters to the given range.

Parameters:float – Values are clipped between -threshold, threshold
clip_gradient_norm(self, threshold: float)

Clip the norm of the gradient by the given value.

Parameters:float – Threshold to clip at
class flambe.nn.SoftmaxLayer(input_size: int, output_size: int, take_log: bool = True)[source]

Bases: flambe.nn.module.Module

Implement an SoftmaxLayer module.

Can be used to form a classifier out of any encoder.

forward(self, data: Tensor)

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – The input data, as a float tensor
Returns:The encoded output, as a float tensor
Return type:torch.Tensor
class flambe.nn.MixtureOfSoftmax(input_size: int, output_size: int, k: int = 1, take_log: bool = True)[source]

Bases: flambe.nn.module.Module

Implement the MixtureOfSoftmax output layer.

pi

softmax layer over the different softmax

Type:FullyConnected
layers

list of the k softmax layers

Type:[FullyConnected]
forward(self, data: Tensor)

Implement mixture of softmax for language modeling.

Parameters:data (torch.Tensor) – seq_len x batch_size x hidden_size
Returns:out – output matrix of shape seq_len x batch_size x out_size
Return type:Variable
class flambe.nn.Embeddings[source]

Bases: flambe.nn.module.Module, torch.nn.Embedding

Implement an Embedding module.

This object replicates the usage of nn.Embedding but registers the from_pretrained classmethod to be used inside a Flambé configuration, as this does not happen automatically during the registration of PyTorch objects.

classmethod from_pretrained(cls, embeddings: Tensor, freeze: bool = True, paddinx_idx: Optional[int] = None, max_norm: Optional[float] = None, norm_type: float = 2.0, scale_grad_by_freq: bool = False, sparse: bool = False)

Create Embedding instance from given 2-dimensional Tensor.

Parameters:
  • embeddings (torch.Tensor) – FloatTensor containing weights for the Embedding. First dimension is being passed to Embedding as num_embeddings, second as embedding_dim.
  • freeze (bool) – If True, the tensor does not get updated in the learning process. Default: True
  • (int, optional) (padding_idx) – See module initialization documentation.
  • max_norm (float, optional) – See module initialization documentation.
  • norm_type (float, optional) – See module initialization documentation. Default 2.
  • scale_grad_by_freq (bool, optional) – See module initialization documentation. Default False.
  • (bool, optional) (sparse) – See module initialization documentation. Default False.
class flambe.nn.Embedder(embedding: nn.Embedding, encoder: Module, embedding_dropout: float = 0, pad_index: Optional[int] = 0)[source]

Bases: flambe.nn.module.Module

Implements an Embedder module.

An Embedder takes as input a sequence of index tokens, and computes the corresponding embedded representations, and padding mask. The encoder may be initialized using a pretrained embedding matrix.

embeddings

The embedding layer

Type:Embedding
encoder

The sub-encoder that this object is wrapping

Type:Encoder
drop

The dropout layer

Type:nn.Dropout
forward(self, data: Tensor)

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – The input data, as a float tensor, batch first
Returns:The encoded output, as a float tensor. May return a state if the encoder is an RNN
Return type:Union[Tensor, Tuple[Tensor, Tensor]]
class flambe.nn.MLPEncoder(input_size: int, output_size: int, n_layers: int = 1, output_activation: Optional[nn.Module] = None, hidden_size: Optional[int] = None, hidden_activation: Optional[nn.Module] = None)[source]

Bases: flambe.nn.module.Module

Implements a multi layer feed forward network.

This module can be used to create output layers, or more complex multi-layer feed forward networks.

seq

the sequence of layers and activations

Type:nn.Sequential
forward(self, data: torch.Tensor)

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – input to the model of shape (batch_size, input_size)
Returns:output – output of the model of shape (batch_size, output_size)
Return type:torch.Tensor
class flambe.nn.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)[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 delt with by using torch’s PackedSequence.

rnn

The rnn submodule

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

Performs a forward pass through the network.

Parameters:data (Tensor) – The input data, as a float tensor
Returns:
  • Tensor – The encoded output, as a float tensor
  • Tensor – The encoded state, as a float tensor
class flambe.nn.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, mask: Optional[Tensor] = None)

Perform a forward pass through the network.

Parameters:data (torch.Tensor) – The input data, as a float tensor
Returns:The encoded output, as a float tensor
Return type:torch.Tensor
class flambe.nn.CNNEncoder(input_channels: int, channels: List[int], conv_dim: int = 2, kernel_size: Union[int, List[Union[Tuple[int, ...], int]]] = 3, activation: nn.Module = None, pooling: nn.Module = None, dropout: float = 0, batch_norm: bool = True, stride: int = 1, padding: int = 0)[source]

Bases: flambe.nn.module.Module

Implements a multi-layer n-dimensional CNN.

This module can be used to create multi-layer CNN models.

cnn

The cnn submodule

Type:nn.Module
forward(self, data: Tensor)

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – The input data, as a float tensor
Returns:The encoded output, as a float tensor
Return type:Union[Tensor, Tuple[Tensor, ..]]
class flambe.nn.Sequential(**kwargs: Dict[str, Union[Module, torch.nn.Module]])[source]

Bases: flambe.nn.Module

Implement a Sequential module.

This class can be used in the same way as torch’s nn.Sequential, with the difference that it accepts kwargs arguments.

forward(self, data: torch.Tensor)

Performs a forward pass through the network.

Parameters:data (torch.Tensor) – input to the model
Returns:output – output of the model
Return type:torch.Tensor