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

Cannot save pruned model with MultiHeadAttention Layer

Open
#1,077 2 comments 0 reactions 1 assignee View on GitHub

Nobody has claimed this yet.

Assessment

This issue has not been assessed yet.

Description

bug

Describe the bug
Trying to save a model that wraps a MultiHeadAttention layer in a PruneLowMagnitude, fails with duplicate dataset name.

System information

TensorFlow version (installed from source or binary): 2.13.0rc1

TensorFlow Model Optimization version (installed from source or binary): 0.7.5

Python version: 3.10

Describe the expected behavior
Successful model save.

Describe the current behavior
When saving a pruned model, I get a ValueError: Unable to create dataset (name already exists) on "mask:0".

Code to reproduce the issue

import tensorflow as tf
import tensorflow_model_optimization as tfmot
import tempfile

if __name__ == '__main__':
    # model
    inputs = tf.keras.layers.Input(shape=(28, 28, 3))
    x = tf.keras.layers.Conv2D(filters=128, kernel_size=3, activation='relu')(inputs)
    x = tf.keras.layers.MultiHeadAttention(num_heads=4, key_dim=128)(query=x, value=x, key=x)
    outputs = tf.keras.layers.Flatten()(x)
    model = tf.keras.Model(inputs=inputs, outputs=outputs)

    model.compile(optimizer='adam', loss='mse')

    # call model to initialize weights
    model(tf.ones((1, 28, 28, 3)))

    # prune model
    pruning_params = {
        'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(
            initial_sparsity=0.5,
            final_sparsity=0.9,
            begin_step=0,
            end_step=1,
            frequency=1,
        ),
    }
    model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, **pruning_params)

    with tempfile.TemporaryDirectory() as temp_dir:
        model_for_pruning.save(temp_dir + '/model.h5')  # <-- fails

Potentially related to #661 and #944.

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.