Transformer#

class torch_ecg.models.Transformer(input_size: int, hidden_size: int, num_heads: int = 8, num_layers: int = 1, dropout: float = 0.1, batch_first: bool = True, **kwargs: Any)[source]#

Bases: Module, SizeMixin

Transformer feature extractor.

Parameters:
  • input_size (int) – Number of input channels.

  • hidden_size (int) – Number of hidden units in the encoding layer.

  • num_heads (int, default 8) – Number of attention heads.

  • num_layers (int, default 1) – Number of encoding layers.

  • dropout (float, default 0.1) – Dropout probability

  • batch_first (bool, default True) – If True, the input is of shape (batch_size, seq_len, input_size), otherwise, the input is of shape (seq_len, batch_size, input_size).

  • kwargs (dict) – Other hyper-parameters of the Module, including activation, kernel_initializer, etc.

compute_output_shape(seq_len: int | None = None, batch_size: int | None = None) → Sequence[int | None][source]#

Compute the output shape of the transformer feature extractor.

Parameters:
  • seq_len (int, optional) – Length of the input tensor.

  • batch_size (int, optional) – Batch size of the input tensor.

Returns:

output_shape – Output shape of the transformer feature extractor.

Return type:

sequence

forward(x: Tensor) → Tensor[source]#

Forward pass of the transformer feature extractor.

Parameters:

x (torch.Tensor) – Input tensor, of shape (seq_len, batch_size, input_size) if batch_first is False, else (batch_size, seq_len, input_size).

Returns:

Output tensor, of shape (seq_len, batch_size, output_size) if batch_first is False, else (batch_size, seq_len, output_size).

Return type:

torch.Tensor