iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

RNN / LSTM

Recurrent Neural Networks process sequences one step at a time, maintaining a hidden state. SimpleRNN, LSTM, and GRU are the three flavours — LSTM is the default for serious sequence work, GRU for tighter compute budgets, SimpleRNN for short toy examples only.

LSTM, GRU, masking, sequence-to-sequence

EXAMPLE
import tensorflow as tf
import tensorflow.keras as keras
from tensorflow.keras import layers
import numpy as np

# 1) Toy data — character-level next-character prediction
text = (open('shakespeare.txt').read()).lower()
chars = sorted(set(text))
char2idx = {c: i for i, c in enumerate(chars)}
idx2char = np.array(chars)
VOCAB = len(chars)
SEQ_LEN = 100

encoded = np.array([char2idx[c] for c in text])

def make_dataset(seq, batch=64):
    ds = tf.data.Dataset.from_tensor_slices(seq)
    ds = ds.batch(SEQ_LEN + 1, drop_remainder=True)
    ds = ds.map(lambda b: (b[:-1], b[1:]))
    ds = ds.shuffle(10_000).batch(batch).prefetch(tf.data.AUTOTUNE)
    return ds

train_ds = make_dataset(encoded[:90_000])
val_ds   = make_dataset(encoded[90_000:])

# 2) LSTM model — input → embedding → LSTM → dense logits
model = keras.Sequential([
    layers.Embedding(VOCAB, 64),
    layers.LSTM(128, return_sequences=True),
    layers.LSTM(128, return_sequences=True),
    layers.Dense(VOCAB),
])
model.compile(
    optimizer='adam',
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['sparse_categorical_accuracy'],
)
model.summary()

# 3) Train
model.fit(
    train_ds, validation_data=val_ds, epochs=20,
    callbacks=[
        keras.callbacks.EarlyStopping(patience=3, restore_best_weights=True),
        keras.callbacks.ModelCheckpoint('best.keras', save_best_only=True),
    ],
)

# 4) GRU — fewer params, often as good
gru = keras.Sequential([
    layers.Embedding(VOCAB, 64),
    layers.GRU(128, return_sequences=True),
    layers.Dense(VOCAB),
])
gru.compile(optimizer='adam',
                       loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True))

# 5) Bidirectional — read forward + backward
bi = keras.Sequential([
    layers.Embedding(VOCAB, 64),
    layers.Bidirectional(layers.LSTM(128, return_sequences=True)),
    layers.Dense(VOCAB),
])
# Only valid when the WHOLE sequence is available (offline tagging, not generation).

# 6) Stateful RNNs — hidden state persists across batches
stateful = keras.Sequential([
    layers.Embedding(VOCAB, 64, batch_input_shape=(1, None)),
    layers.LSTM(128, return_sequences=True, stateful=True),
    layers.Dense(VOCAB),
])
# Use for generation: reset_states() between sequences; pass one token at a time.

# 7) Generate text from a trained model
def sample(model, seed='hello', n=200, temperature=1.0):
    ids = [char2idx[c] for c in seed]
    model.reset_states()
    inp = tf.expand_dims(ids, 0)
    out = list(seed)
    for _ in range(n):
        logits = model(inp)[:, -1, :] / temperature
        idx = tf.random.categorical(logits, num_samples=1)[0, 0].numpy()
        out.append(idx2char[idx])
        inp = tf.expand_dims([idx], 0)
    return ''.join(out)

# 8) Variable-length sequences — masking
# Most production data has variable-length sequences. Mask the padding.
model = keras.Sequential([
    layers.Embedding(VOCAB, 64, mask_zero=True),         # PAD token = 0; mask propagates through layers
    layers.LSTM(128, return_sequences=True),
    layers.Dense(VOCAB),
])

# Padded batch
x = tf.keras.preprocessing.sequence.pad_sequences([
    [1, 2, 3, 4, 5],
    [6, 7, 8, 0, 0],
], padding='post')
# mask propagates: layers ignore positions where input == 0.

# 9) Sequence labeling — many-to-many (NER, POS)
from tensorflow.keras.layers import TimeDistributed

ner = keras.Sequential([
    layers.Embedding(VOCAB, 64, mask_zero=True),
    layers.Bidirectional(layers.LSTM(128, return_sequences=True)),
    layers.TimeDistributed(layers.Dense(NUM_TAGS, activation='softmax')),
])
ner.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# 10) Sequence classification — many-to-one (sentiment, intent)
cls = keras.Sequential([
    layers.Embedding(VOCAB, 64, mask_zero=True),
    layers.Bidirectional(layers.LSTM(64)),                # last hidden state
    layers.Dense(64, activation='relu'),
    layers.Dense(1,  activation='sigmoid'),
])
cls.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# 11) Seq2Seq — encoder + decoder for translation, summarisation
# (Modern practice: transformers. But LSTMs still useful for short sequences + low memory.)

encoder_inputs = keras.Input(shape=(None,))
x = layers.Embedding(VOCAB, 64)(encoder_inputs)
_, state_h, state_c = layers.LSTM(128, return_state=True)(x)
enc_states = [state_h, state_c]

decoder_inputs = keras.Input(shape=(None,))
y = layers.Embedding(VOCAB, 64)(decoder_inputs)
out, _, _ = layers.LSTM(128, return_sequences=True, return_state=True)(y, initial_state=enc_states)
outputs = layers.Dense(VOCAB, activation='softmax')(out)

seq2seq = keras.Model([encoder_inputs, decoder_inputs], outputs)
seq2seq.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# 12) Attention (BahdanauAttention / LuongAttention)
import tensorflow_addons as tfa
# Or use keras.layers.Attention / MultiHeadAttention for transformer-style attention

attention = layers.MultiHeadAttention(num_heads=4, key_dim=32)
# Use INSTEAD of plain LSTM stacks for long-range dependencies. Most modern NLP uses transformers.

# 13) Performance tips
# • Pad to a SMALL multiple of 8 / 16 — better GPU utilisation
# • Use mixed_precision for 2-3x speedup on supported GPUs:
#       keras.mixed_precision.set_global_policy('mixed_float16')
# • Set dropout via the layer arg: layers.LSTM(128, dropout=0.2, recurrent_dropout=0.0)
#       recurrent_dropout disables CuDNN kernel — much slower; leave at 0 for speed
# • CuDNN kernel requires: activation='tanh', recurrent_activation='sigmoid', use_bias=True, no recurrent_dropout, no projections
# • Profile with tf.profiler / TensorBoard to find bottlenecks (often dataset, not the GPU)

# 14) When to use transformers instead
# • Long sequences (> 200 tokens)
# • You can afford their compute
# • You have lots of labelled data OR can fine-tune a pretrained model
# Transformers (BERT, T5, GPT) dominate NLP benchmarks for a reason; LSTMs make sense when
# you need a tiny model, on-device inference, or true streaming.

# 15) Common bugs
# • Forgetting return_sequences=True between stacked RNN layers → second layer sees only the last step
# • Mask propagation broken — apply mask_zero=True on the FIRST embedding only; subsequent layers respect it
# • Stateful RNN without reset_states between sequences → state bleeds across unrelated sequences
# • Using a stateful RNN with batch_size != 1 in inference — must match training batch shape
# • Generation with from_logits=False but feeding raw logits — use softmax or set from_logits in loss
# • Long sequences + small batch size + recurrent_dropout > 0 — terrible throughput
# • Forgetting time dimension in TimeDistributed wrapper — Dense applies per timestep correctly only when wrapped

Why it matters

Reach for LSTM (or GRU for tighter budgets) for short-to-medium sequences and tasks where you want a small model on-device. For long-range dependencies and modern NLP, transformers win — but understanding RNNs is still the cleanest entry point to sequence modelling, and the CuDNN-optimised kernels remain genuinely fast for typical inference.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
from tensorflow.keras import layers, Sequential
model = Sequential([
    layers.Input((None, 64)),
    layers.LSTM(128),
    layers.Dense(1),
])
Try it Yourself »

Discussion

Loading…