DepthwiseConv2D Layers cannot be clustered or sparsely pruned

[nathanwbailey]

When attempting to sparsely prune or cluster DepthwiseConv2D Layers it appears that no clustering or sparse pruning actually occurs.

System information

TensorFlow version (installed from source or binary): 2.16.1/2.15.1

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

Python version: 3.9.18

Describe the expected behavior

I should see that the kernel for the Depthwise layer should have 3 unique weights, not 9.

Describe the current behavior

The Depthwise layer has 9 unique weights

The same issue occurs for sparse m by n pruning, the weights are not pruned correctly in an m by n manner

Code to reproduce the issue

import tensorflow as tf
import tensorflow_model_optimization as tfmot
from tensorflow import keras

import numpy as np

(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()
train_images = train_images / 255.0
test_images = test_images / 255.0

model = keras.Sequential([
keras.layers.InputLayer(input_shape=(28, 28)),
keras.layers.Reshape(target_shape=(28, 28, 1)),
keras.layers.DepthwiseConv2D(kernel_size=(3, 3), activation=tf.nn.relu),
keras.layers.MaxPooling2D(pool_size=(2, 2)),
keras.layers.Flatten(),
keras.layers.Dense(10)
])

model.compile(optimizer='adam',
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])

model.fit(
train_images,
train_labels,
validation_split=0.1,
epochs=10
)

cluster_weights = tfmot.clustering.keras.cluster_weights
CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

clustering_params = {
'number_of_clusters': 3,
'cluster_centroids_init': CentroidInitialization.LINEAR
}

clustered_model = cluster_weights(model, **clustering_params)

opt = keras.optimizers.Adam(learning_rate=1e-5)

clustered_model.compile(
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=opt,
metrics=['accuracy'])
clustered_model.fit(
train_images,
train_labels,
batch_size=500,
epochs=1,
validation_split=0.1)

for layer in model.layers:
for weight in layer.weights:
print(weight.name)
print(len(np.unique(weight)))