Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

float16 quantization runs out of memory for LSTM model

Open
#1,091 3 comments 0 reactions 1 assignee View on GitHub

@yyoon is already working on this.

Since Aug 31, 2023.

Assessment

This issue has not been assessed yet.

Description

bug

No matter the size of the LSTM model, converting it with float16 optimization runs out of memory.

Code to reproduce the issue
The code snippet to reproduce the issue on Google Colab
Code:

import numpy as np
import tensorflow as tf
import tensorflow_model_optimization as tfmot

def create_model():
  model = tf.keras.models.Sequential()

  # For the model to later get converted, batch_size and sequence_length should be fixed.
  # E.g., batch_input_shape=[None, 1] will throw an error.
  # This is just a limitation when using RNNs. E.g., for FC or CNN we can have batch_size=None
  model.add(tf.keras.layers.Embedding(
    input_dim=5,
    output_dim=1,
    batch_input_shape=[1, 1]
  ))

  model.add(tf.keras.layers.LSTM(
    units=1,
    return_sequences=False,
    stateful=False
  ))

  model.add(tf.keras.layers.Dense(5))

  return model

model = create_model()
model.summary()

model.save("/content/model/")

representative_data = np.random.randint(0, 5, (200, 1)).astype(np.float32)

def representative_dataset():
  for sample in representative_data:
    sample = np.expand_dims(sample, axis=0)     # batch_size = 1
    yield [sample]                              # set sample as first (and only) input of the model

# float16 quantization
converter = tf.lite.TFLiteConverter.from_saved_model("/content/model/")
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]
# kernel runs out of memory and crashes in the following line
tflite_quant_model = converter.convert()
Dominant language
Python
Stars
1.6k
Forks
349
Avg merge
3d 2h
Merged PRs (30d)
1

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from tensorflow/model-optimization

All issues in tensorflow/model-optimization

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.