rtnn.models package

Submodules

rtnn.models.rnn module

class rtnn.models.rnn.BaseRNN(*args: Any, **kwargs: Any)[source]

Bases: Module

Base class for bidirectional RNN modules (LSTM/GRU).

This class provides a common interface for both LSTM and GRU models with bidirectional processing and a final 1D convolutional layer to map the hidden states to the desired output channels.

Parameters:

feature_channel (int) – Number of input features per time step.
output_channel (int) – Number of output channels (target variables).
hidden_size (int) – Number of hidden units in the RNN layers.
num_layers (int) – Number of stacked RNN layers.
rnn_type (str) – Type of RNN cell, either ‘lstm’ or ‘gru’.

rnn

The bidirectional RNN layer.

Type:: nn.LSTM or nn.GRU

final

Final 1D convolution to project hidden states to output channels.

Type:: nn.Conv1d

hidden_size

Number of hidden units.

Type:: int

num_layers

Number of stacked layers.

Type:: int

output_channel

Number of output channels.

Type:: int

Examples

>>> model = BaseRNN(
...     feature_channel=6,
...     output_channel=4,
...     hidden_size=64,
...     num_layers=2,
...     rnn_type='lstm'
... )
>>> x = torch.randn(32, 6, 10)  # (batch, features, sequence)
>>> y = model(x)
>>> y.shape
torch.Size([32, 4, 10])

__init__(feature_channel: int, output_channel: int, hidden_size: int, num_layers: int, rnn_type: str) → None[source]

Initialize the BaseRNN module.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
hidden_size (int) – Size of hidden state.
num_layers (int) – Number of RNN layers.
rnn_type (str) – Type of RNN (‘lstm’ or ‘gru’).

Raises:

ValueError – If rnn_type is not ‘lstm’ or ‘gru’.

init_hidden(batch_size: int, device: torch.device) → torch.Tensor | Tuple[torch.Tensor, torch.Tensor][source]

Initialize the hidden state for the RNN.

Parameters:

batch_size (int) – Batch size for the input.
device (torch.device) – Device to create the hidden state on.

Returns:

For GRU: returns hidden state tensor of shape (2 * num_layers, batch_size, hidden_size) For LSTM: returns tuple (hidden, cell) both of same shape.

Return type:

torch.Tensor or tuple of torch.Tensor

forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass through the bidirectional RNN.

Parameters:: x (torch.Tensor) – Input tensor of shape (batch_size, feature_channel, seq_length).
Returns:: Output tensor of shape (batch_size, output_channel, seq_length).
Return type:: torch.Tensor

Notes

The input is permuted to (batch_size, seq_length, feature_channel) for the RNN, then the output is permuted back for the convolution.

class rtnn.models.rnn.RNN_LSTM(*args: Any, **kwargs: Any)[source]

Bases: BaseRNN

LSTM-based bidirectional RNN model.

This class inherits from BaseRNN and configures it to use LSTM cells.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
hidden_size (int) – Size of hidden state.
num_layers (int) – Number of LSTM layers.

Examples

>>> model = RNN_LSTM(
...     feature_channel=6,
...     output_channel=4,
...     hidden_size=128,
...     num_layers=3
... )
>>> x = torch.randn(16, 6, 10)
>>> y = model(x)
>>> print(y.shape)
torch.Size([16, 4, 10])

__init__(feature_channel: int, output_channel: int, hidden_size: int, num_layers: int) → None[source]

Initialize the LSTM model.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
hidden_size (int) – Size of hidden state.
num_layers (int) – Number of LSTM layers.

class rtnn.models.rnn.RNN_GRU(*args: Any, **kwargs: Any)[source]

Bases: BaseRNN

GRU-based bidirectional RNN model.

This class inherits from BaseRNN and configures it to use GRU cells.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
hidden_size (int) – Size of hidden state.
num_layers (int) – Number of GRU layers.

Examples

>>> model = RNN_GRU(
...     feature_channel=6,
...     output_channel=4,
...     hidden_size=128,
...     num_layers=3
... )
>>> x = torch.randn(16, 6, 10)
>>> y = model(x)
>>> print(y.shape)
torch.Size([16, 4, 10])

__init__(feature_channel: int, output_channel: int, hidden_size: int, num_layers: int) → None[source]

Initialize the GRU model.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
hidden_size (int) – Size of hidden state.
num_layers (int) – Number of GRU layers.

rtnn.models.fcn module

class rtnn.models.fcn.FCBlock(*args: Any, **kwargs: Any)[source]

Bases: Module

A fully connected block with linear layer, batch normalization, and ReLU activation.

This module applies a linear transformation, followed by batch normalization, and then a ReLU activation function.

Parameters:

in_features (int) – Number of input features.
out_features (int) – Number of output features.

linear

Linear transformation layer.

Type:: nn.Linear

bn

Batch normalization layer.

Type:: nn.BatchNorm1d

relu

ReLU activation function.

Type:: nn.ReLU

Examples

>>> block = FCBlock(128, 64)
>>> x = torch.randn(32, 128)
>>> y = block(x)
>>> y.shape
torch.Size([32, 64])

__init__(in_features: int, out_features: int) → None[source]

Initialize the FCBlock.

Parameters:

in_features (int) – Number of input features.
out_features (int) – Number of output features.

forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass through the FCBlock.

Parameters:: x (torch.Tensor) – Input tensor of shape (batch_size, in_features).
Returns:: Output tensor of shape (batch_size, out_features).
Return type:: torch.Tensor

Notes

The forward pass applies: ReLU(BatchNorm(Linear(x)))

class rtnn.models.fcn.FCN(*args: Any, **kwargs: Any)[source]

Bases: Module

Fully Connected Network with configurable depth and width.

This model flattens the input sequence and processes it through a series of fully connected layers. It can optionally expand the sequence length using a linear transformation.

Parameters:

feature_channel (int) – Number of input features per time step.
output_channel (int) – Number of output channels.
num_layers (int) – Number of hidden layers.
hidden_size (int) – Size of hidden layers.
seq_length (int, optional) – Length of the input sequence. Default is 55.
dim_expand (int, optional) – Number of time steps to expand the output sequence by. Default is 0 (no expansion).

feature_channel

Number of input features.

Type:: int

output_channel

Number of output channels.

Type:: int

seq_length

Length of the input sequence.

Type:: int

dim_expand

Number of time steps to expand by.

Type:: int

input_layer

First fully connected layer.

Type:: FCBlock

hidden_layers

Stack of hidden layers.

Type:: nn.Sequential

output_layer

Final output layer.

Type:: nn.Linear

dim_change

Optional layer for sequence length expansion.

Type:: nn.Linear or None

Examples

>>> model = FCN(
...     feature_channel=6,
...     output_channel=4,
...     num_layers=3,
...     hidden_size=196,
...     seq_length=10
... )
>>> x = torch.randn(32, 6, 10)
>>> y = model(x)
>>> y.shape
torch.Size([32, 4, 10])

__init__(feature_channel: int, output_channel: int, num_layers: int, hidden_size: int, seq_length: int = 55, dim_expand: int = 0) → None[source]

Initialize the FCN model.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
num_layers (int) – Number of hidden layers.
hidden_size (int) – Size of hidden layers.
seq_length (int, optional) – Length of the input sequence. Default is 55.
dim_expand (int, optional) – Number of time steps to expand the output sequence by. Default is 0 (no expansion).

Raises:

ValueError – If num_layers is less than 1.

forward(x: torch.Tensor) → torch.Tensor[source]

Forward pass through the FCN.

Parameters:: x (torch.Tensor) – Input tensor of shape (batch_size, feature_channel, seq_length).
Returns:: Output tensor of shape (batch_size, output_channel, seq_length + dim_expand) if dim_expand > 0, otherwise (batch_size, output_channel, seq_length).
Return type:: torch.Tensor

Notes

The forward pass: 1. Flattens the input to (batch_size, feature_channel * seq_length) 2. Passes through FCBlocks 3. Projects to output dimensions 4. Reshapes to (batch_size, output_channel, seq_length) 5. Optionally expands sequence length

rtnn.models.Transformer module

class rtnn.models.Transformer.SelfAttention(*args: Any, **kwargs: Any)[source]

Bases: Module

Multi-head self-attention mechanism.

This module implements scaled dot-product attention with multiple heads.

Parameters:

embed_size (int) – Size of the embedding dimension.
heads (int) – Number of attention heads.

embed_size

Embedding dimension.

Type:: int

heads

Number of attention heads.

Type:: int

head_dim

Dimension of each attention head (embed_size // heads).

Type:: int

values

Linear layer for value projections.

Type:: nn.Linear

keys

Linear layer for key projections.

Type:: nn.Linear

queries

Linear layer for query projections.

Type:: nn.Linear

fc_out

Final output linear layer.

Type:: nn.Linear

Examples

>>> attention = SelfAttention(embed_size=128, heads=4)
>>> values = torch.randn(32, 10, 128)
>>> keys = torch.randn(32, 10, 128)
>>> query = torch.randn(32, 10, 128)
>>> out = attention(values, keys, query, mask=None)
>>> out.shape
torch.Size([32, 10, 128])

__init__(embed_size: int, heads: int) → None[source]

Initialize the SelfAttention module.

Parameters:

embed_size (int) – Size of the embedding dimension.
heads (int) – Number of attention heads.

Raises:

AssertionError – If embed_size is not divisible by heads.

forward(values: torch.Tensor, keys: torch.Tensor, query: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor[source]

Forward pass through the self-attention mechanism.

Parameters:

values (torch.Tensor) – Value tensor of shape (batch_size, value_len, embed_size).
keys (torch.Tensor) – Key tensor of shape (batch_size, key_len, embed_size).
query (torch.Tensor) – Query tensor of shape (batch_size, query_len, embed_size).
mask (torch.Tensor, optional) – Attention mask of shape (batch_size, 1, 1, key_len) or (batch_size, query_len, key_len). Default is None.

Returns:

Output tensor of shape (batch_size, query_len, embed_size).

Return type:

torch.Tensor

Notes

The attention mechanism follows the formula: Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) V

class rtnn.models.Transformer.TransformerBlock(*args: Any, **kwargs: Any)[source]

Bases: Module

Transformer block consisting of self-attention and feed-forward layers.

This module applies multi-head self-attention followed by a feed-forward network, with layer normalization and residual connections.

Parameters:

embed_size (int) – Size of the embedding dimension.
heads (int) – Number of attention heads.
dropout (float) – Dropout rate.
forward_expansion (int) – Expansion factor for the feed-forward network.

attention

Multi-head self-attention module.

Type:: SelfAttention

norm1

First layer normalization.

Type:: nn.LayerNorm

norm2

Second layer normalization.

Type:: nn.LayerNorm

feed_forward

Feed-forward network.

Type:: nn.Sequential

dropout

Dropout layer.

Type:: nn.Dropout

Examples

>>> block = TransformerBlock(128, 4, dropout=0.1, forward_expansion=4)
>>> x = torch.randn(32, 10, 128)
>>> out = block(x, x, x, mask=None)
>>> out.shape
torch.Size([32, 10, 128])

__init__(embed_size: int, heads: int, dropout: float, forward_expansion: int) → None[source]

Initialize the TransformerBlock.

Parameters:

embed_size (int) – Size of the embedding dimension.
heads (int) – Number of attention heads.
dropout (float) – Dropout rate.
forward_expansion (int) – Expansion factor for the feed-forward network.

forward(value: torch.Tensor, key: torch.Tensor, query: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor[source]

Forward pass through the transformer block.

Parameters:

value (torch.Tensor) – Value tensor of shape (batch_size, seq_len, embed_size).
key (torch.Tensor) – Key tensor of shape (batch_size, seq_len, embed_size).
query (torch.Tensor) – Query tensor of shape (batch_size, seq_len, embed_size).
mask (torch.Tensor, optional) – Attention mask. Default is None.

Returns:

Output tensor of shape (batch_size, seq_len, embed_size).

Return type:

torch.Tensor

class rtnn.models.Transformer.Encoder(*args: Any, **kwargs: Any)[source]

Bases: Module

Transformer encoder for sequence-to-sequence processing.

This module applies positional encoding and a stack of transformer blocks to transform input sequences.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
embed_size (int) – Size of the embedding dimension.
num_layers (int) – Number of transformer blocks.
heads (int) – Number of attention heads.
forward_expansion (int) – Expansion factor for feed-forward networks.
seq_length (int) – Length of the input sequence.
dropout (float) – Dropout rate.

embed_size

Embedding dimension.

Type:: int

seq_length

Input sequence length.

Type:: int

first

Initial linear projection.

Type:: nn.Linear

first_act

Activation function.

Type:: nn.ReLU

position_embedding

Positional embeddings.

Type:: nn.Embedding

layers

Stack of transformer blocks.

Type:: nn.ModuleList

dropout

Dropout layer.

Type:: nn.Dropout

final

Final convolution to map to output channels.

Type:: nn.Conv1d

Examples

>>> encoder = Encoder(
...     feature_channel=6,
...     output_channel=4,
...     embed_size=64,
...     num_layers=2,
...     heads=4,
...     forward_expansion=4,
...     seq_length=10,
...     dropout=0.1
... )
>>> x = torch.randn(32, 6, 10)
>>> y = encoder(x)
>>> y.shape
torch.Size([32, 4, 10])

__init__(feature_channel: int, output_channel: int, embed_size: int, num_layers: int, heads: int, forward_expansion: int, seq_length: int, dropout: float) → None[source]

Initialize the Transformer encoder.

Parameters:

feature_channel (int) – Number of input features.
output_channel (int) – Number of output channels.
embed_size (int) – Size of the embedding dimension.
num_layers (int) – Number of transformer blocks.
heads (int) – Number of attention heads.
forward_expansion (int) – Expansion factor for feed-forward networks.
seq_length (int) – Length of the input sequence.
dropout (float) – Dropout rate.

Raises:

ValueError – If num_layers is less than 1.

forward(x: torch.Tensor, mask: torch.Tensor | None = None) → torch.Tensor[source]

Forward pass through the transformer encoder.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, feature_channel, seq_length).
mask (torch.Tensor, optional) – Attention mask. Default is None.

Returns:

Output tensor of shape (batch_size, output_channel, seq_length).

Return type:

torch.Tensor

Notes

The forward pass: 1. Permutes input to (batch, seq, features) 2. Applies linear projection to embeddings 3. Adds positional embeddings 4. Passes through transformer blocks 5. Permutes back and applies final convolution

class rtnn.models.Transformer.EncoderTorch(*args: Any, **kwargs: Any)[source]

Bases: Module

__init__(feature_channel: int, output_channel: int, embed_size: int, num_layers: int, heads: int, forward_expansion: int, seq_length: int, dropout: float) → None[source]

forward(x: torch.Tensor, mask: torch.Tensor | None = None, src_key_padding_mask: torch.Tensor | None = None) → torch.Tensor[source]: x: (batch, feature_channel, seq_length)

rtnn.models.DimChangeModule module

class rtnn.models.DimChangeModule.DimChange(*args: Any, **kwargs: Any)[source]

Bases: Module

B*V*H -> B*V*(H + H_add)

__init__(channel_number, output_number)[source]

forward(x)[source]