def testClusterKerasUnsupportedLayer(self):
"""Verifies that attempting to cluster an unsupported layer raises an exception."""
keras_unsupported_layer = self.keras_unsupported_layer
# We need to build weights before check.
keras_unsupported_layer.build(input_shape=(10, 10))
with self.assertRaises(ValueError):
cluster.cluster_weights(keras_unsupported_layer, **self.params)
def testClusterKerasUnsupportedLayer(self):
"""
Verifies that attempting to cluster an unsupported layer raises an
exception.
"""
with self.assertRaises(ValueError):
cluster.cluster_weights(self.keras_unsupported_layer, **self.params)
def testClusterModelUnsupportedKerasLayerRaisesError(self):
with self.assertRaises(ValueError):
cluster.cluster_weights(
keras.Sequential([
self.keras_clusterable_layer,
self.keras_non_clusterable_layer,
self.custom_clusterable_layer, self.keras_unsupported_layer
]), **self.params)
def testClusterModelUnsupportedKerasLayerRaisesError(self):
"""Verifies that attempting to cluster a model that contains an unsupported layer raises an exception."""
keras_unsupported_layer = self.keras_unsupported_layer
# We need to build weights before check.
keras_unsupported_layer.build(input_shape=(10, 10))
with self.assertRaises(ValueError):
cluster.cluster_weights(
keras.Sequential([
self.keras_clusterable_layer, self.keras_non_clusterable_layer,
self.custom_clusterable_layer, keras_unsupported_layer
]), **self.params)
def testClusterModelCustomNonClusterableLayerRaisesError(self):
"""Verifies that attempting to cluster a model that contains a custom non-clusterable layer raises an exception."""
with self.assertRaises(ValueError):
custom_non_clusterable_layer = self.custom_non_clusterable_layer
# We need to build weights before check.
custom_non_clusterable_layer.build(input_shape=(1, 2))
cluster.cluster_weights(
keras.Sequential([
self.keras_clusterable_layer, self.keras_non_clusterable_layer,
self.custom_clusterable_layer, custom_non_clusterable_layer
]), **self.params)
def testClusterModelCustomNonClusterableLayerRaisesError(self):
"""
Verifies that attempting to cluster a model that contains a custom
non-clusterable layer raises an exception.
"""
with self.assertRaises(ValueError):
cluster.cluster_weights(
keras.Sequential([
self.keras_clusterable_layer, self.keras_non_clusterable_layer,
self.custom_clusterable_layer, self.custom_non_clusterable_layer
]), **self.params)
def testClusterModelDoesNotWrapAlreadyWrappedLayer(self):
model = keras.Sequential([
layers.Flatten(),
cluster.cluster_weights(layers.Dense(10), **self.params),
])
clustered_model = cluster.cluster_weights(model, **self.params)
clustered_model.build(input_shape=(10, 10, 1))
self.assertEqual(len(model.layers), len(clustered_model.layers))
self._validate_clustered_layer(model.layers[0],
clustered_model.layers[0])
# Second layer is used as-is since it's already a clustered layer.
self.assertEqual(model.layers[1], clustered_model.layers[1])
self._validate_clustered_layer(model.layers[1].layer,
clustered_model.layers[1])
def testEndToEndClusterPreserve(self):
"""Runs CQAT end to end and whole model is quantized."""
original_model = tf.keras.Sequential(
[layers.Dense(5, activation='softmax', input_shape=(10, ))])
clustered_model = cluster.cluster_weights(original_model,
**self.cluster_params)
self.compile_and_fit(clustered_model)
clustered_model = cluster.strip_clustering(clustered_model)
num_of_unique_weights_clustering = self._get_number_of_unique_weights(
clustered_model, 0, 'kernel')
quant_aware_annotate_model = (
quantize.quantize_annotate_model(clustered_model))
quant_aware_model = quantize.quantize_apply(
quant_aware_annotate_model,
scheme=default_8bit_cluster_preserve_quantize_scheme.
Default8BitClusterPreserveQuantizeScheme())
self.compile_and_fit(quant_aware_model)
stripped_cqat_model = strip_clustering_cqat(quant_aware_model)
# Check the unique weights of a certain layer of
# clustered_model and pcqat_model
num_of_unique_weights_cqat = self._get_number_of_unique_weights(
stripped_cqat_model, 1, 'kernel')
self.assertAllEqual(num_of_unique_weights_clustering,
num_of_unique_weights_cqat)
def cluster_train_eval_strip(
model, x_train, y_train, x_test, y_test, batch_size, test_case):
"""Train, evaluate and strip clustering."""
model = cluster.cluster_weights(
model,
number_of_clusters=16,
cluster_centroids_init=cluster_config.CentroidInitialization
.KMEANS_PLUS_PLUS,)
model.compile(loss="binary_crossentropy",
optimizer="adam",
metrics=["accuracy"])
print("Train...")
model.fit(x_train, y_train, batch_size=batch_size, epochs=1,
validation_data=(x_test, y_test), verbose=2)
score, acc = model.evaluate(x_test, y_test,
batch_size=batch_size)
print("Test score:", score)
print("Test accuracy:", acc)
print("Strip clustering wrapper...")
model = cluster.strip_clustering(model)
if "Bidirectional" in test_case:
layer_weight = getattr(model.layers[1].forward_layer.cell, "kernel")
elif "StackedRNNCells" in test_case:
layer_weight = getattr(model.layers[1].cell.cells[0], "kernel")
else:
raise ValueError("Only Bidirectional and StackedRNNCells are tested now.")
print("Number of clusters:", len(set(layer_weight.numpy().flatten())))
def cluster_model(model, x_train, y_train, x_test, y_test):
print('Clustering model')
clustering_params = {
'number_of_clusters':
8,
'cluster_centroids_init':
cluster_config.CentroidInitialization.DENSITY_BASED
}
# Cluster model
clustered_model = cluster.cluster_weights(model, **clustering_params)
# Use smaller learning rate for fine-tuning
# clustered model
opt = tf.keras.optimizers.Adam(learning_rate=1e-5)
clustered_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=opt,
metrics=['accuracy'])
# Fine-tune clustered model
clustered_model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs_fine_tuning,
verbose=1,
validation_split=0.1)
score = clustered_model.evaluate(x_test, y_test, verbose=0)
print('Clustered model test loss:', score[0])
print('Clustered model test accuracy:', score[1])
return clustered_model
def testValuesAreClusteredAfterStripping(self, number_of_clusters,
cluster_centroids_init):
"""
Verifies that, for any number of clusters and any centroid initialization
method, the number of unique weight values after stripping is always less
or equal to number_of_clusters.
"""
original_model = tf.keras.Sequential([
layers.Dense(32, input_shape=(10, )),
])
self.assertGreater(
len(set(original_model.get_weights()[0].reshape(-1, ).tolist())),
number_of_clusters)
clustered_model = cluster.cluster_weights(
original_model,
number_of_clusters=number_of_clusters,
cluster_centroids_init=cluster_centroids_init)
stripped_model = cluster.strip_clustering(clustered_model)
weights_as_list = stripped_model.get_weights()[0].reshape(
-1, ).tolist()
unique_weights = set(weights_as_list)
# Make sure numbers match
self.assertLessEqual(len(unique_weights), number_of_clusters)
# Make sure that the stripped layer is the Dense one
self.assertIsInstance(stripped_model.layers[0], layers.Dense)
def _cluster_model(model, number_of_clusters):
(x_train, y_train), _ = _get_dataset()
clustering_params = {
'number_of_clusters':
number_of_clusters,
'cluster_centroids_init':
cluster_config.CentroidInitialization.KMEANS_PLUS_PLUS
}
# Cluster model
clustered_model = cluster.cluster_weights(model, **clustering_params)
# Use smaller learning rate for fine-tuning
# clustered model
opt = tf.keras.optimizers.Adam(learning_rate=1e-5)
clustered_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=opt,
metrics=['accuracy'])
# Fine-tune clustered model
clustered_model.fit(x_train, y_train, epochs=EPOCHS_FINE_TUNING)
stripped_model = cluster.strip_clustering(clustered_model)
stripped_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer=opt,
metrics=['accuracy'])
return stripped_model
def testEndToEnd(self):
"""Test End to End clustering."""
original_model = keras.Sequential([
layers.Dense(2, input_shape=(2,)),
layers.Dense(2),
])
clustered_model = cluster.cluster_weights(original_model, **self.params)
clustered_model.compile(
loss=keras.losses.categorical_crossentropy,
optimizer="adam",
metrics=["accuracy"],
)
clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
stripped_model = cluster.strip_clustering(clustered_model)
_, tflite_file = tempfile.mkstemp(".tflite")
_, keras_file = tempfile.mkstemp(".h5")
if not compat.is_v1_apis():
converter = tf.lite.TFLiteConverter.from_keras_model(stripped_model)
else:
tf.keras.models.save_model(stripped_model, keras_file)
converter = tf.lite.TFLiteConverter.from_keras_model_file(keras_file)
tflite_model = converter.convert()
with open(tflite_file, "wb") as f:
f.write(tflite_model)
self._verify_tflite(tflite_file, self.x_train)
os.remove(keras_file)
os.remove(tflite_file)
def testClusterModelDoesNotWrapAlreadyWrappedLayer(self):
"""Verifies that clustering a model that contains an already clustered layer does not result in wrapping the clustered layer into another cluster_wrapper."""
model = keras.Sequential(
[
layers.Flatten(),
cluster.cluster_weights(layers.Dense(10), **self.params),
])
clustered_model = cluster.cluster_weights(model, **self.params)
clustered_model.build(input_shape=(10, 10, 1))
self.assertEqual(len(model.layers), len(clustered_model.layers))
self._validate_clustered_layer(model.layers[0], clustered_model.layers[0])
# Second layer is used as-is since it's already a clustered layer.
self.assertEqual(model.layers[1], clustered_model.layers[1])
self._validate_clustered_layer(model.layers[1].layer,
clustered_model.layers[1])
def end_to_end_testing(self, original_model, clusters_check=None):
"""Test End to End clustering."""
clustered_model = cluster.cluster_weights(original_model,
**self.params)
clustered_model.compile(
loss=keras.losses.categorical_crossentropy,
optimizer="adam",
metrics=["accuracy"],
)
clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=1)
stripped_model = cluster.strip_clustering(clustered_model)
if clusters_check is not None:
clusters_check(stripped_model)
_, tflite_file = tempfile.mkstemp(".tflite")
_, keras_file = tempfile.mkstemp(".h5")
converter = tf.lite.TFLiteConverter.from_keras_model(stripped_model)
tflite_model = converter.convert()
with open(tflite_file, "wb") as f:
f.write(tflite_model)
self._verify_tflite(tflite_file, self.x_test)
os.remove(keras_file)
os.remove(tflite_file)
def testClusterSubclassModel(self):
"""
Verifies that attempting to cluster an instance of a subclass of
keras.Model raises an exception.
"""
model = TestModel()
with self.assertRaises(ValueError):
_ = cluster.cluster_weights(model, **self.params)
def testClusterSequentialModelWithInput(self):
# With InputLayer
model = keras.Sequential([
layers.Dense(10, input_shape=(10, )),
layers.Dense(10),
])
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(self._count_clustered_layers(clustered_model), 2)
def _clusterTrainStrip(self, model):
clustered_model = cluster.cluster_weights(
model,
**self.params_clustering,
)
self._train(clustered_model)
stripped_model = cluster.strip_clustering(clustered_model)
return stripped_model
def testClusterSequentialModelWithInput(self):
"""Verifies that a sequential model with an input layer is being clustered correctly."""
# With InputLayer
model = keras.Sequential([
layers.Dense(10, input_shape=(10,)),
layers.Dense(10),
])
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(self._count_clustered_layers(clustered_model), 2)
def testClusterFunctionalModel(self):
i1 = keras.Input(shape=(10, ))
i2 = keras.Input(shape=(10, ))
x1 = layers.Dense(10)(i1)
x2 = layers.Dense(10)(i2)
outputs = layers.Add()([x1, x2])
model = keras.Model(inputs=[i1, i2], outputs=outputs)
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(self._count_clustered_layers(clustered_model), 3)
def testClusterFunctionalModelWithLayerReused(self):
# The model reuses the Dense() layer. Make sure it's only clustered once.
inp = keras.Input(shape=(10, ))
dense_layer = layers.Dense(10)
x = dense_layer(inp)
x = dense_layer(x)
model = keras.Model(inputs=[inp], outputs=[x])
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(self._count_clustered_layers(clustered_model), 1)
def testClusterSubclassModelAsSubmodel(self):
"""
Verifies that attempting to cluster a model with submodel
that is a subclass throws an exception.
"""
model_subclass = TestModel()
model = keras.Sequential([layers.Dense(10), model_subclass])
with self.assertRaisesRegexp(ValueError, "Subclassed models.*"):
_ = cluster.cluster_weights(model, **self.params)
def testClusterValidLayersListSuccessful(self):
model_layers = [
self.keras_clusterable_layer, self.keras_non_clusterable_layer,
self.custom_clusterable_layer
]
clustered_list = cluster.cluster_weights(model_layers, **self.params)
self.assertEqual(len(model_layers), len(clustered_list))
for layer, clustered_layer in zip(model_layers, clustered_list):
self._validate_clustered_layer(layer, clustered_layer)
def testStripClusteringSequentialModel(self):
model = keras.Sequential([
layers.Dense(10),
layers.Dense(10),
])
clustered_model = cluster.cluster_weights(model, **self.params)
stripped_model = cluster.strip_clustering(clustered_model)
self.assertEqual(self._count_clustered_layers(stripped_model), 0)
self.assertEqual(model.get_config(), stripped_model.get_config())
def testClusterValidLayersListSuccessful(self):
"""Verifies that clustering a list of layers results in all clusterable layers within the list being clustered."""
model_layers = [
self.keras_clusterable_layer, self.keras_non_clusterable_layer,
self.custom_clusterable_layer
]
clustered_list = cluster.cluster_weights(model_layers, **self.params)
self.assertEqual(len(model_layers), len(clustered_list))
for layer, clustered_layer in zip(model_layers, clustered_list):
self._validate_clustered_layer(layer, clustered_layer)
def testStripSelectivelyClusteredSequentialModel(self):
clustered_model = keras.Sequential([
cluster.cluster_weights(layers.Dense(10), **self.params),
layers.Dense(10),
])
clustered_model.build(input_shape=(1, 10))
stripped_model = cluster.strip_clustering(clustered_model)
self.assertEqual(self._count_clustered_layers(stripped_model), 0)
self.assertIsInstance(stripped_model.layers[0], layers.Dense)
def testStripClusteringSequentialModel(self):
"""Verifies that stripping the clustering wrappers from a sequential model produces the expected config."""
model = keras.Sequential([
layers.Dense(10),
layers.Dense(10),
])
clustered_model = cluster.cluster_weights(model, **self.params)
stripped_model = cluster.strip_clustering(clustered_model)
self.assertEqual(self._count_clustered_layers(stripped_model), 0)
self.assertEqual(model.get_config(), stripped_model.get_config())
def testStripSelectivelyClusteredSequentialModel(self):
"""Verifies that invoking strip_clustering() on a selectively clustered sequential model strips the clustering wrappers from the clustered layers."""
clustered_model = keras.Sequential([
cluster.cluster_weights(layers.Dense(10), **self.params),
layers.Dense(10),
])
clustered_model.build(input_shape=(1, 10))
stripped_model = cluster.strip_clustering(clustered_model)
self.assertEqual(self._count_clustered_layers(stripped_model), 0)
self.assertIsInstance(stripped_model.layers[0], layers.Dense)
def testStripSelectivelyClusteredFunctionalModel(self):
i1 = keras.Input(shape=(10, ))
i2 = keras.Input(shape=(10, ))
x1 = cluster.cluster_weights(layers.Dense(10), **self.params)(i1)
x2 = layers.Dense(10)(i2)
outputs = layers.Add()([x1, x2])
clustered_model = keras.Model(inputs=[i1, i2], outputs=outputs)
stripped_model = cluster.strip_clustering(clustered_model)
self.assertEqual(self._count_clustered_layers(stripped_model), 0)
self.assertIsInstance(stripped_model.layers[2], layers.Dense)
def testClusterSequentialModelSelectively(self):
"""
Verifies that layers within a sequential model can be clustered
selectively.
"""
clustered_model = keras.Sequential()
clustered_model.add(cluster.cluster_weights(self.keras_clusterable_layer, **self.params))
clustered_model.add(self.keras_clusterable_layer)
clustered_model.build(input_shape=(1, 10))
self.assertIsInstance(clustered_model.layers[0], cluster_wrapper.ClusterWeights)
self.assertNotIsInstance(clustered_model.layers[1], cluster_wrapper.ClusterWeights)
