{"id":22447,"library":"tensorflow-model-optimization","title":"TensorFlow Model Optimization Toolkit","description":"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.","status":"active","version":"0.8.0","language":"python","source_language":"en","source_url":"https://github.com/tensorflow/model-optimization","tags":["tensorflow","model-optimization","quantization","pruning","clustering","deployment"],"install":[{"cmd":"pip install tensorflow-model-optimization","lang":"bash","label":"Install from PyPI"}],"dependencies":[{"reason":"Core dependency; the toolkit works with TensorFlow 2.x.","package":"tensorflow","optional":false}],"imports":[{"note":"Direct import may lead to submodule not found errors; use tfmot.quantization.keras.quantize_model","wrong":"from tensorflow_model_optimization import quantize_model","symbol":"quantize_model","correct":"import tensorflow_model_optimization as tfmot"},{"note":"Incorrect alias; use full package path","wrong":"from tfmot.sparsity import prune_low_magnitude","symbol":"prune_low_magnitude","correct":"from tensorflow_model_optimization.sparsity.keras import prune_low_magnitude"},{"note":"Standard alias, no common mistake observed.","wrong":null,"symbol":"tfmot","correct":"import tensorflow_model_optimization as tfmot"}],"quickstart":{"code":"import tensorflow as tf\nimport tensorflow_model_optimization as tfmot\n\n# Quantization aware training example\ndef get_model():\n    model = tf.keras.Sequential([\n        tf.keras.layers.Dense(32, activation='relu', input_shape=(784,)),\n        tf.keras.layers.Dense(10, activation='softmax')\n    ])\n    return model\n\nmodel = get_model()\nquantize_model = tfmot.quantization.keras.quantize_model\nq_aware_model = quantize_model(model)\n\n# Compile and train\nq_aware_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])\nimport numpy as np\ntrain_images = np.random.random((100, 784)).astype('float32')\ntrain_labels = np.random.randint(10, size=(100,))\nq_aware_model.fit(train_images, train_labels, epochs=1, verbose=0)\n\n# Convert to TFLite\nconverter = tf.lite.TFLiteConverter.from_keras_model(q_aware_model)\nconverter.optimizations = [tf.lite.Optimize.DEFAULT]\ntflite_model = converter.convert()\nprint('TFLite model size:', len(tflite_model))","lang":"python","description":"Demonstrates quantization-aware training with a simple model, then converts to TFLite."},"warnings":[{"fix":"Ensure you are using tf.keras, not keras directly. Upgrade to TensorFlow 2.14+.","message":"The 'quantize_model' function only works with Keras v2 (tf.keras). Using tfmot with tf.compat.v1.keras or standalone Keras may fail.","severity":"gotcha","affected_versions":">=0.8.0"},{"fix":"Re-apply quantization after upgrading to 0.7.0+. For backward compatibility, use tfmot.quantization.keras.QuantizeWrapper (old) instead of QuantizeWrapperV2 if needed.","message":"In 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.","severity":"breaking","affected_versions":">=0.7.0"},{"fix":"Avoid decorating quantized/pruned models with @tf.function. Instead, call the model eagerly or use Keras training loop.","message":"Pruning and QAT layers cannot be used with tf.function tracing when combined. Debugging is tricky.","severity":"gotcha","affected_versions":">=0.6.0"},{"fix":"Use QAT or pruning instead; consider alternatives like TensorFlow Lite's built-in clustering.","message":"Weight clustering API is no longer actively maintained and may be removed in future versions.","severity":"deprecated","affected_versions":">=0.8.0"},{"fix":"Replace PeepholeLSTMCell with standard LSTM or custom implementation before upgrading.","message":"Version 0.7.2 removed support for PeepholeLSTMCell. Loading models using this layer will fail.","severity":"breaking","affected_versions":">=0.7.2"}],"env_vars":null,"last_verified":"2026-04-27T00:00:00.000Z","next_check":"2026-07-26T00:00:00.000Z","problems":[{"fix":"Use: from tensorflow_model_optimization.quantization.keras import quantize_model","cause":"The submodule may not be imported correctly; the function is inside a deeper package.","error":"ModuleNotFoundError: No module named 'tensorflow_model_optimization.quantization.keras.quantize_model'"},{"fix":"Use: with tfmot.quantization.keras.quantize_scope(): model = tf.keras.models.load_model('path')","cause":"When loading a quantized model, the custom QuantizeWrapper class is not registered.","error":"ValueError: Unknown layer: QuantizeWrapper. Please ensure this layer is imported."},{"fix":"Use: from tensorflow_model_optimization.sparsity.keras import prune_low_magnitude","cause":"Importing tfmot.sparsity directly fails because the package is not fully imported.","error":"AttributeError: module 'tensorflow_model_optimization' has no attribute 'sparsity'"}],"ecosystem":"pypi","meta_description":null,"install_score":null,"install_tag":null,"quickstart_score":null,"quickstart_tag":null}