Registry / ai-ml / tensorflow-model-optimization

tensorflow-model-optimization

library0.8.1pypypi✓ verified 86d ago

A suite of tools for optimizing machine learning models for deployment, including quantization-aware training, pruning, and weight clustering. Current version 0.8.0, tested against TensorFlow 2.14.1. Released approximately every 4–6 months.

ai-ml llm-agents

pip install tensorflow-model-optimization

github ↗package ↗

INSTALL

IMPORT

SIG · TENSORFLOW-MODEL-O

tensorflow-model-optimization

ai-mlpythonv0.8.1

Install

36.9s avg

Import

13520ms

Disk

2355MB

Pass rate

5/ 10

Env Coverage5 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v0.8.1 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

py 3.10–3.95 runs

build_error

glibc

py 3.10–3.95 runs

installs and imports cleanly · install 36.9s · import 13.520s · 2355.2MB

2355MB installed

● package 2355MB

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

quantize_model

✓ import tensorflow_model_optimization as tfmot

✗ from tensorflow_model_optimization import quantize_model

Direct import may lead to submodule not found errors; use tfmot.quantization.keras.quantize_model

prune_low_magnitude

✓ from tensorflow_model_optimization.sparsity.keras import prune_low_magnitude

✗ from tfmot.sparsity import prune_low_magnitude

Incorrect alias; use full package path

tfmot

✓ import tensorflow_model_optimization as tfmot

Standard alias, no common mistake observed.

Demonstrates quantization-aware training with a simple model, then converts to TFLite.

import tensorflow as tf import tensorflow_model_optimization as tfmot # Quantization aware training example def get_model(): model = tf.keras.Sequential([ tf.keras.layers.Dense(32, activation='relu', input_shape=(784,)), tf.keras.layers.Dense(10, activation='softmax') ]) return model model = get_model() quantize_model = tfmot.quantization.keras.quantize_model q_aware_model = quantize_model(model) # Compile and train q_aware_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy']) import numpy as np train_images = np.random.random((100, 784)).astype('float32') train_labels = np.random.randint(10, size=(100,)) q_aware_model.fit(train_images, train_labels, epochs=1, verbose=0) # Convert to TFLite converter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model) converter.optimizations = [tf.lite.Optimize.DEFAULT] tflite_model = converter.convert() print('TFLite model size:', len(tflite_model))

Debug

Known issues

gotchaThe 'quantize_model' function only works with Keras v2 (tf.keras). Using tfmot with tf.compat.v1.keras or standalone Keras may fail.

fix

Ensure you are using tf.keras, not keras directly. Upgrade to TensorFlow 2.14+.

affects: >=0.8.0

breakingIn version 0.7.0+, the default QAT API changed: QuantizeWrapperV2 now preserves the order of weights. Existing models serialized with QuantizeWrapper may not load correctly.

fix

Re-apply quantization after upgrading to 0.7.0+. For backward compatibility, use tfmot.quantization.keras.QuantizeWrapper (old) instead of QuantizeWrapperV2 if needed.

affects: >=0.7.0

gotchaPruning and QAT layers cannot be used with tf.function tracing when combined. Debugging is tricky.

fix

Avoid decorating quantized/pruned models with @tf.function. Instead, call the model eagerly or use Keras training loop.

affects: >=0.6.0

deprecatedWeight clustering API is no longer actively maintained and may be removed in future versions.

fix

Use QAT or pruning instead; consider alternatives like TensorFlow Lite's built-in clustering.

affects: >=0.8.0

breakingVersion 0.7.2 removed support for PeepholeLSTMCell. Loading models using this layer will fail.

fix

Replace PeepholeLSTMCell with standard LSTM or custom implementation before upgrading.

affects: >=0.7.2

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'tensorflow_model_optimization.quantization.keras.quantize_model'

The submodule may not be imported correctly; the function is inside a deeper package.

fix

Use: from tensorflow_model_optimization.quantization.keras import quantize_model

ValueError: Unknown layer: QuantizeWrapper. Please ensure this layer is imported.

When loading a quantized model, the custom QuantizeWrapper class is not registered.

fix

Use: with tfmot.quantization.keras.quantize_scope(): model = tf.keras.models.load_model('path')

AttributeError: module 'tensorflow_model_optimization' has no attribute 'sparsity'

Importing tfmot.sparsity directly fails because the package is not fully imported.

fix

Use: from tensorflow_model_optimization.sparsity.keras import prune_low_magnitude

Upgrade

Version history

0.8.1latest on PyPI · released May 12, 2026

Audit

Dependencies

tensorflowrequiredCore dependency; the toolkit works with TensorFlow 2.x.

Agent activity

41 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/tensorflow/model-optimization ↗

packagehttps://pypi.org/project/tensorflow-model-optimization/ ↗