Skip to content

tsmixer

TsMixer Model

Overview

TsMixer is a fully MLP-based architecture for time series data.

For more info, refer to the original paper TsMixer: An All-MLP Architecture for Time Series.

Classes:

Functions:

Classes

TsMixerBlockParams

TsMixer block parameters

Attributes:

  • norm (Literal['batch', 'layer']) –

    Normalization type

  • activation (Literal['relu', 'gelu']) –

    Activation type

  • dropout (float) –

    Dropout rate

  • ff_dim (int) –

    Feed forward dimension

TsMixerParams

TsMixer parameters

Attributes:

  • blocks (list[TsBlockParams]) –

    TsMixer blocks

  • name (str) –

    Model name

TsMixerModel

Helper class to generate model from parameters

Functions

layer_from_params staticmethod 
layer_from_params(inputs: keras.Input, params: TsMixerParams | dict, num_classes: int | None = None)

Create layer from parameters

Source code in neuralspot_edge/models/tsmixer.py 
147
148
149
150
151
152
@staticmethod
def layer_from_params(inputs: keras.Input, params: TsMixerParams | dict, num_classes: int | None = None):
    """Create layer from parameters"""
    if isinstance(params, dict):
        params = TsMixerParams(**params)
    return tsmixer_layer(x=inputs, params=params, num_classes=num_classes)
model_from_params staticmethod 
model_from_params(inputs: keras.Input, params: TsMixerParams | dict, num_classes: int | None = None)

Create model from parameters

Source code in neuralspot_edge/models/tsmixer.py 
154
155
156
157
158
@staticmethod
def model_from_params(inputs: keras.Input, params: TsMixerParams | dict, num_classes: int | None = None):
    """Create model from parameters"""
    outputs = TsMixerModel.layer_from_params(inputs=inputs, params=params, num_classes=num_classes)
    return keras.Model(inputs=inputs, outputs=outputs)

Functions

norm_layer 

norm_layer(norm: str, name: str) -> keras.Layer

Normalization layer

Parameters:

  • norm (str) –

    Normalization type

  • name (str) –

    Name

Returns:

  • Layer –

    keras.Layer: Layer

Source code in neuralspot_edge/models/tsmixer.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def norm_layer(norm: str, name: str) -> keras.Layer:
    """Normalization layer

    Args:
        norm (str): Normalization type
        name (str): Name

    Returns:
        keras.Layer: Layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        """Functional normalization layer

        Args:
            x (keras.KerasTensor): Input tensor

        Returns:
            keras.KerasTensor: Output tensor
        """
        if norm == "batch":
            return keras.layers.BatchNormalization(axis=[-2, -1], name=f"{name}.BN")(x)
        if norm == "layer":
            return keras.layers.LayerNormalization(axis=[-2, -1], name=f"{name}.LN")(x)
        return x

    return layer

ts_block 

ts_block(params: TsMixerBlockParams, name: str) -> keras.Layer

Residual block of TSMixer.

Parameters:

  • params (TsBlockParams) –

    Block parameters

  • name (str) –

    Name

Returns:

  • Layer –

    keras.Layer: Layer

Source code in neuralspot_edge/models/tsmixer.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def ts_block(params: TsMixerBlockParams, name: str) -> keras.Layer:
    """Residual block of TSMixer.

    Args:
        params (TsBlockParams): Block parameters
        name (str): Name

    Returns:
        keras.Layer: Layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        # Temporal Linear
        y = norm_layer(params.norm, name=f"{name}.TL")(x)

        y = keras.ops.transpose(y, axes=[0, 2, 1])  # [Batch, Channel, Input Length]
        y = keras.layers.Dense(y.shape[-1], activation=params.activation, name=f"{name}.TL.DENSE")(y)
        y = keras.ops.transpose(y, axes=[0, 2, 1])  # [Batch, Input Length, Channel]
        y = keras.layers.Dropout(params.dropout, name=f"{name}.TL.DROP")(y)
        res = y + x

        # Feature Linear
        y = norm_layer(params.norm, name=f"{name}.FL")(res)
        y = keras.layers.Dense(params.ff_dim, activation=params.activation, name=f"{name}.FL.DENSE")(
            y
        )  # [Batch, Input Length, FF_Dim]
        y = keras.layers.Dropout(params.dropout, name=f"{name}.FL.DROP")(y)

        y = keras.layers.Dense(x.shape[-1], name=f"{name}.RL.DENSE")(y)  # [Batch, Input Length, Channel]
        y = keras.layers.Dropout(params.dropout, name=f"{name}.RL.DROP")(y)
        return y + res

    return layer

tsmixer_layer 

tsmixer_layer(inputs: keras.KerasTensor, params: any, num_classes: int) -> keras.KerasTensor

TsMixer layer

Parameters:

  • inputs (KerasTensor) –

    Input tensor

  • params (any) –

    Model parameters

  • num_classes (int) –

    Number of classes

Returns:

  • KerasTensor –

    keras.KerasTensor: Output tensor

Source code in neuralspot_edge/models/tsmixer.py 
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def tsmixer_layer(inputs: keras.KerasTensor, params: any, num_classes: int) -> keras.KerasTensor:
    """TsMixer layer

    Args:
        inputs (keras.KerasTensor): Input tensor
        params (any): Model parameters
        num_classes (int): Number of classes

    Returns:
        keras.KerasTensor: Output tensor
    """
    y = inputs
    for block in range(params.blocks):
        y = ts_block(block)(y)

    # if target_slice:
    #     y = y[:, :, target_slice]

    y = keras.ops.transpose(y, axes=[0, 2, 1])  # [Batch, Channel, Input Length]
    y = keras.layers.Dense(num_classes)(y)  # [Batch, Channel, Output Length]
    y = keras.ops.transpose(y, axes=[0, 2, 1])  # [Batch, Output Length, Channel])

    return y