flambe.nn.pooling

Module Contents

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