Train bidirectional-LSTM

!pip install git+https://github.com/ASEM000/serket --quiet
!pip install optax --quiet

Imports

import jax
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import optax  # for gradient optimization
import serket as sk
import functools as ft
from typing_extensions import Annotated
import time

EPOCHS = 100
LR = 1e-3
Input = Annotated[jax.Array, "Float[seq_len, input_dim]"]

Data preparation

# 101 samples of 1D data
x = jnp.linspace(0, 1, 101).reshape(-1, 1)
y = jnp.sin(2 * jnp.pi * x)
y += 0.1 * jr.normal(jr.key(0), y.shape)
# we will use 2 time steps to predict the next time step
x_train = jnp.stack([x[:-1], x[1:]], axis=1)
# 100 minibatches x 1 sample x 2 time steps x 1D data
x_train = jnp.reshape(x_train, (100, 1, 2, 1))
# 100 minibatches x 1 samples x 1D data
y_train = jnp.reshape(y[1:], (100, 1, 1))

Model creation

class BiLstm(sk.TreeClass):
    def __init__(
        self,
        in_features: int,
        hidden_dim: int,
        out_features: int,
        *,
        key: jax.Array,
    ):
        k1, k2, k3 = jr.split(key, 3)
        self.cell1 = sk.nn.LSTMCell(in_features, hidden_dim, key=k1)
        self.cell2 = sk.nn.LSTMCell(in_features, hidden_dim, key=k2)
        self.cell3 = sk.nn.LSTMCell(hidden_dim * 2, out_features, key=k3)

    def __call__(self, input: Input) -> Input:
        # initialize the states of the cells
        state = sk.tree_state(self)
        # run the forward cell
        output1, state1 = sk.nn.scan_cell(self.cell1)(input, state.cell1)
        # run the backward cell
        output2, state2 = sk.nn.scan_cell(self.cell2, reverse=True)(input, state.cell2)
        # concatenate the outputs
        output = jnp.concatenate((output1, output2), axis=1)
        # run the final cell
        output, state3 = sk.nn.scan_cell(self.cell3)(output, state.cell3)
        # return the last time step
        return output[-1]


key = jax.random.key(0)
net = BiLstm(1, 64, 1, key=key)
# 1) mask the non-jaxtype parameters
net = sk.tree_mask(net)
# 2) initialize the optimizer state
optim = optax.adam(LR)
optim_state = optim.init(net)

Training functions

def mse(x, y):
    assert x.shape == y.shape
    return jnp.mean((x - y) ** 2)


@ft.partial(jax.grad, has_aux=True)
def loss_func(net: BiLstm, x: jax.Array, y: jax.Array):
    # pass non-jaxtype over jax transformation
    # using `tree_mask`/`tree_unmask` scheme
    # 3) unmask the non-jaxtype parameters to be used in the computation
    net = sk.tree_unmask(net)
    # 4) vectorize the computation over the batch dimension
    # and get the logits
    # here we dont vectorize over state argument so we use `None`
    logits = jax.vmap(net)(x)
    # 5) use the appropriate loss function
    loss = mse(logits, y)
    return loss, (loss, logits)


@jax.jit
def train_step(net: BiLstm, optim_state: optax.OptState, x: jax.Array, y: jax.Array):
    grads, (loss, logits) = loss_func(net, x, y)
    updates, optim_state = optim.update(grads, optim_state)
    net = optax.apply_updates(net, updates)
    return net, optim_state, (loss, logits)

Train and plot results

for i in range(1, EPOCHS + 1):
    t0 = time.time()
    for j, (xb, yb) in enumerate(zip(x_train, y_train)):
        net, optim_state, (loss, logits) = train_step(net, optim_state, xb, yb)
        print(
            f"Epoch: {i:003d}/{EPOCHS:003d}\t"
            f"Batch: {j+1:003d}/{len(x_train):003d}\t"
            f"Batch loss: {loss:3e}\t"
            f"Time: {time.time() - t0:.3f}",
            end="\r",
        )


# 6) un-mask the trained network
eval_net = sk.tree_unmask(net)


y_pred = jax.vmap(eval_net)(x_train.reshape(-1, 2, 1))
plt.plot(x[1:], y[1:], "--k", label="data")
plt.plot(x[1:], y_pred, label="prediction")
plt.legend()

Epoch: 100/100  Batch: 100/100  Batch loss: 2.065103e-03        Time: 0.022

<matplotlib.legend.Legend at 0x298c1dc10>