LSTM from Pytorch to Flax NNX / Performance issues LSTM

[JavierJJR]

Hi 👋,
I'm new to using Jax and I'm finding it challenging to convert my models from PyTorch to Flax NNX, particularly the LSTMs. For example, the following model processes data in batches of length 64 with 553 timesteps and 1 feature (64, 553, 1). How would the same model be implemented in Flax NNX? I feel that it would be helpful to supplement the package's documentation with more examples to assist new users.

class RNN_Model(nn.Module):
    def __init__(self, input_size=1, hidden_size=128, num_layers=2, output_size=10):
        super(RNN_Model, self).__init__()

        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size,
                            num_layers=num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        lstm_out, _ = self.lstm(x) 
        out = self.fc(lstm_out[:, -1, :])  
        return out

[Tushaam]

I can suggest a possible implementation of this in JAX/Flax however i am not 100% certain that this Flax code is working exactly as the pytorch implementation but lets give it a try

so to implement this:

1. since in JAX we need to code certain things manually, we need to be careful for tracking each operation especially for a complex NN like LSTM

so in PyTorch that is done in 2 lines..In JAX we need to be alert regarding the input and how it is being passed to the LSTM cell and then track each hidden state,cell state

so as u know for each time instance t we would need to compute the hidden state h and cell state c
since u have 2 layers in ur LSTM net, so we would do it something like this:

carry = [
OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(i), (batch_size,), self.hidden_size)
for i in range(self.num_layers)
]

but we also need to loop over time instances which is:
for t in range(x.shape[1]):
input_t = x[:, t, :]
for l in range(self.num_layers):
carry[l], input_t = OptimizedLSTMCell()(carry[l], input_t)

so as per me the full code should be:

class LSTM(nn.Module):
hidden_size: int = 128
output_size: int = 10
num_layers: int = 2

@nn.compact
def __call__(self, x):  
    batch_size = x.shape[0]

    


    carry = [
        OptimizedLSTMCell.initialize_carry(jax.random.PRNGKey(i), (batch_size,), self.hidden_size)
        for i in range(self.num_layers)
    ]

    
    for t in range(x.shape[1]):
        input_t = x[:, t, :]
        for l in range(self.num_layers):
            carry[l], input_t = OptimizedLSTMCell()(carry[l], input_t)

if anyone thinks am missing something or got something wrong, please let me know!!
And @JavierJJR i hope this was of some help to you😀

[JavierJJR]

Hi @Tushaam ,
First of all, thank you so much for your response — it was really helpful. After giving it some thought, I’ve come to the conclusion that using nnx.scan is the best approach (I am using Flax NNX).

However, I’ve implemented the same LSTM network using both Flax and PyTorch, and surprisingly, I’ve found that the Flax version is significantly slower than the PyTorch one.

Since I’m still new to the JAX ecosystem, I’m not entirely sure whether there’s something wrong with how I defined the model or the training loop that could be causing such a noticeable slowdown.

I’m attaching the code I used to compare the models in Flax and PyTorch — hopefully someone can help identify the reason behind the much longer training times in Flax. Note that I am running this code on Colab's GPU T4.

Flax Code

#Create synthetic database (sine wave) and DATALOADER

import numpy as np
import tensorflow as tf

tf.random.set_seed(0)
np.random.seed(42)

train_steps = 1200
eval_every = 200
batch_size = 64
sequence_length = 553
n_features = 1

def generate_sine_wave_data(n_samples, seq_length, n_features):
    data = np.zeros((n_samples, seq_length, n_features))
    labels = np.zeros((n_samples, 1))
    
    for i in range(n_samples):
        # Random frequency and phase
        freq = np.random.uniform(0.1, 0.5)
        phase = np.random.uniform(0, 2*np.pi)
        
        t = np.arange(seq_length)
        sine_wave = np.sin(2 * np.pi * freq * t + phase)
        
        noise = np.random.normal(0, 0.1, seq_length)
        signal = sine_wave + noise
        
        data[i, :, 0] = signal
       
        labels[i, 0] = freq
    
    return data, labels

train_data, train_labels = generate_sine_wave_data(8000, sequence_length, n_features)
test_data, test_labels = generate_sine_wave_data(2000, sequence_length, n_features)

train_ds = tf.data.Dataset.from_tensor_slices((train_data, train_labels))
test_ds = tf.data.Dataset.from_tensor_slices((test_data, test_labels))

train_ds = train_ds.repeat().shuffle(1024)
train_ds = train_ds.batch(batch_size, drop_remainder=True).take(train_steps).prefetch(1)
test_ds = test_ds.batch(batch_size, drop_remainder=True).prefetch(1)

for x_batch, y_batch in train_ds.take(1):
    print("Input batch shape:", x_batch.shape)
    print("Label batch shape:", y_batch.shape)

import jax.numpy as jnp
import optax
from flax import nnx
import jax

class LSTM(nnx.Module):
  def __init__(self, input_size: int, hidden_size: int, rngs: nnx.Rngs):
    self.hidden_size = hidden_size
    self.cell1 = nnx.nn.recurrent.OptimizedLSTMCell(input_size, self.hidden_size, rngs=rngs)
    self.cell2 = nnx.nn.recurrent.OptimizedLSTMCell(self.hidden_size, self.hidden_size, rngs=rngs)
    self.fc = nnx.Linear(hidden_size, 1, rngs=rngs)

  def __call__(self, x):

    scan_fn1 = lambda carry, cell, x: cell(carry, x)
    carry1 = self.cell1.initialize_carry((x.shape[0], self.hidden_size), rngs=None)
    carry1, y1 = nnx.scan(
      scan_fn1, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)
    )(carry1, self.cell1, x)

    scan_fn2 = lambda carry, cell, x: cell(carry, x)
    carry2 = self.cell2.initialize_carry((x.shape[0], self.hidden_size), rngs=None)
    carry2, y2 = nnx.scan(
      scan_fn2, in_axes=(nnx.Carry, None, 1), out_axes=(nnx.Carry, 1)
    )(carry2, self.cell2, y1)

    y = y2[:, -1, :]
    y = self.fc(y)

    return y

x = jnp.ones((64,553,1))
model = LSTM(x.shape[2], 128, rngs=nnx.Rngs(0))
nnx.display(model)

import optax

learning_rate = 0.001
momentum = 0.9

optimizer = nnx.Optimizer(model, optax.adamw(learning_rate, momentum))
metrics = nnx.MultiMetric(
  loss=nnx.metrics.Average('loss'),
)

@nnx.jit  
def train_step(model, optimizer, metrics, x, y):
  def loss_fn(model):
    y_pred = model(x)  
    return ((y_pred - y) ** 2).mean()

  loss, grads = nnx.value_and_grad(loss_fn)(model)
  optimizer.update(grads)  
  metrics.update(loss=loss)

  return loss

cached_train_step = nnx.cached_partial(train_step, model, optimizer, metrics)

# Train

metrics_history = {
  'train_loss': [],
}

for step, batch in enumerate(train_ds.as_numpy_iterator()):
  # Run the optimization for one step and make a stateful update to the following:
  # - The train state's model parameters
  # - The optimizer state
  # - The training loss and accuracy batch metrics
  #train_step(model, optimizer, metrics, batch[0],batch[1])
  cached_train_step( batch[0],batch[1])

  if step > 0 and (step % eval_every == 0 or step == train_steps - 1):  # One training epoch has passed.
    print(step)
    # Log the training metrics.
    for metric, value in metrics.compute().items():  # Compute the metrics.
      metrics_history[f'train_{metric}'].append(value)  # Record the metrics.
    metrics.reset()  # Reset the metrics for the test set.

PyTorch

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# Set random seed for reproducibility
torch.manual_seed(0)
np.random.seed(42)

# Parameters
train_steps = 1200
eval_every = 200
batch_size = 64
sequence_length = 553
n_features = 1

# Generate synthetic data for LSTM
def generate_sine_wave_data(n_samples, seq_length, n_features):
    data = np.zeros((n_samples, seq_length, n_features))
    labels = np.zeros((n_samples, 1))
    
    for i in range(n_samples):
        # Random frequency and phase
        freq = np.random.uniform(0.1, 0.5)
        phase = np.random.uniform(0, 2*np.pi)
        
        # Generate sine wave
        t = np.arange(seq_length)
        sine_wave = np.sin(2 * np.pi * freq * t + phase)
        
        # Add noise
        noise = np.random.normal(0, 0.1, seq_length)
        signal = sine_wave + noise
        
        # Reshape to include feature dimension
        data[i, :, 0] = signal
        
        # Create a label (e.g., frequency class or regression target)
        labels[i, 0] = freq
    
    return data, labels

# Generate training and test data
train_data, train_labels = generate_sine_wave_data(8000, sequence_length, n_features)
test_data, test_labels = generate_sine_wave_data(2000, sequence_length, n_features)

# Create PyTorch Dataset
class SineWaveDataset(Dataset):
    def __init__(self, data, labels):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.labels = torch.tensor(labels, dtype=torch.float32)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# Create DataLoader
train_dataset = SineWaveDataset(train_data, train_labels)
test_dataset = SineWaveDataset(test_data, test_labels)

train_loader = DataLoader(
    train_dataset, 
    batch_size=batch_size,
    shuffle=True,
    drop_last=True
)

test_loader = DataLoader(
    test_dataset,
    batch_size=batch_size,
    shuffle=False,
    drop_last=True
)

# Check the shape of a batch
for x_batch, y_batch in train_loader:
    print("Input batch shape:", x_batch.shape)
    print("Label batch shape:", y_batch.shape)
    break

# Define the LSTM model
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        
        # Two LSTM layers
        self.lstm1 = nn.LSTM(input_size, hidden_size, batch_first=True)
        self.lstm2 = nn.LSTM(hidden_size, hidden_size, batch_first=True)
        
        # Fully connected layer
        self.fc = nn.Linear(hidden_size, 1)
        
    def forward(self, x):
        # First LSTM layer
        lstm1_out, _ = self.lstm1(x)
        
        # Second LSTM layer
        lstm2_out, _ = self.lstm2(lstm1_out)
        
        # Get the last time step output
        y = lstm2_out[:, -1, :]
        
        # Final linear layer
        y = self.fc(y)
        
        return y

# Initialize model, loss function, and optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = LSTMModel(n_features, 128).to(device)
criterion = nn.MSELoss()
optimizer = optim.AdamW(model.parameters(), lr=0.001, betas=(0.9, 0.999))

# Training function
def train(model, train_loader, criterion, optimizer, num_steps, eval_every):
    model.train()
    metrics_history = {'train_loss': []}
    step_count = 0
    
    train_iterator = iter(train_loader)
    
    for step in range(num_steps):
        
        try:
            inputs, targets = next(train_iterator)
        except StopIteration:
            train_iterator = iter(train_loader)
            inputs, targets = next(train_iterator)
        
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        # Forward pass
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        step_count += 1
        
        # Log metrics every eval_every steps
        if step > 0 and (step % eval_every == 0 or step == num_steps - 1):
            print(f"Step {step}, Loss: {loss.item():.4f}")
            metrics_history['train_loss'].append(loss.item())
    
    return metrics_history

metrics_history = train(model, train_loader, criterion, optimizer, train_steps, eval_every)

[DiagRisker]

Glad to see you figuring things out!
using tqdm for rough performance assessment :
for b in tqdm(batch,bar_format="{l_bar}{bar:10}{r_bar}",leave=True):
can you provide a time difference between your Pytorch and NNX implementation?

[JavierJJR]

Hello,
In this case, the training time for the Flax model is 72 seconds (using all train steps), while the PyTorch training period is 30 seconds (running the code on Google Colab's T4 GPU). Given my inexperience with JAX, I understand that my implementation of the LSTM model in Flax may be causing some bottleneck that makes the system considerably slower than PyTorch.

[Tushaam]

Hi @JavierJJR i had a look at your reply and your code. One potential reason why JAX takes longer than PyTorch is the usage of @jit

when i first implemented some codes in JAX and Flax i came across a concept by research that when you make your 1st call to the @jit decorator function, it takes longer time to execute as compared to your 2nd, 3rd... n calls

In PyTorch you dont have such issues since..its simply bcoz PyTorch doesnt perform a compilation step before execution but JAX first compiles the code using XLA , then executes.

I feel that might be the reason why its adding up to the overall execution time of your code with JAX compared to PyTorch. Other than this i am unable to point out any potential reasons for it😅

[JavierJJR]

Hi @Tushaam ,

First of all, thank you for your help regarding this issue ☺️.

Secondly, what you mentioned makes a lot of sense. However, I’ve tried running the code responsible for training the model multiple times, and the training times don’t change substantially. I’ve also run some tests with a few databases I have, and Flax’s LSTM models consistently take twice as long as PyTorch’s models. I haven’t been able to pinpoint what the reason for this issue is.

Honestly, my intention was to migrate all the models I regularly use from PyTorch to Flax, but this is turning out to be a major setback 😓.