flambe.nn.pooling

Module Contents

flambe.nn.pooling._default_padding_mask(data: torch.Tensor) → torch.Tensor[source]

Builds a 1s padding mask taking into account initial 2 dimensions of input data.

Parameters:data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
Returns:A padding mask , as a tensor of shape [B x S]
Return type:torch.Tensor
flambe.nn.pooling._sum_with_padding_mask(data: torch.Tensor, padding_mask: torch.Tensor) → torch.Tensor[source]

Applies padding_mask and performs summation over the data

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The result of the summation, as a tensor of shape [B x H]

Return type:

torch.Tensor

class flambe.nn.pooling.FirstPooling[source]

Bases: flambe.nn.Module

Get the last hidden state of a sequence.

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

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor

class flambe.nn.pooling.LastPooling[source]

Bases: flambe.nn.Module

Get the last hidden state of a sequence.

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

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor

class flambe.nn.pooling.SumPooling[source]

Bases: flambe.nn.Module

Get the sum of the hidden state of a sequence.

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

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor

class flambe.nn.pooling.AvgPooling[source]

Bases: flambe.nn.Module

Get the average of the hidden state of a sequence.

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

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • padding_mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor

class flambe.nn.pooling.StructuredSelfAttentivePooling(input_size: int, attention_heads: int = 16, attention_units: Sequence[int] = (300, ), output_activation: Optional[torch.nn.Module] = None, hidden_activation: Optional[torch.nn.Module] = None, is_biased: bool = False, input_dropout: float = 0.0, attention_dropout: float = 0.0)[source]

Bases: flambe.nn.Module

Structured Self Attentive Pooling.

_compute_attention(self, data: torch.Tensor, mask: torch.Tensor)[source]

Computes the attention

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

The attention, as a tensor of shape [B x S x HEADS]

Return type:

torch.Tensor

forward(self, data: torch.Tensor, mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor

class flambe.nn.pooling.GeneralizedPooling(input_size: int, attention_units: Sequence[int] = (300, ), output_activation: Optional[torch.nn.Module] = None, hidden_activation: Optional[torch.nn.Module] = None, is_biased: bool = True, input_dropout: float = 0.0, attention_dropout: float = 0.0)[source]

Bases: flambe.nn.pooling.StructuredSelfAttentivePooling

Self attention pooling.

forward(self, data: torch.Tensor, mask: Optional[torch.Tensor] = None)[source]

Performs a forward pass.

Parameters:
  • data (torch.Tensor) – The input data, as a tensor of shape [B x S x H]
  • mask (torch.Tensor) – The input mask, as a tensor of shape [B X S]
Returns:

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

Return type:

torch.Tensor