def testClusterSequentialModelPreservesBuiltStateWithInput(self):
# With InputLayer
model = keras.Sequential([
layers.Dense(10, input_shape=(10, )),
layers.Dense(10),
])
self.assertEqual(model.built, True)
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(model.built, True)
# Test built state is preserved across serialization
with cluster.cluster_scope():
loaded_model = keras.models.model_from_config(
json.loads(clustered_model.to_json()))
self.assertEqual(loaded_model.built, True)
def testClusterSequentialModelPreservesBuiltStateNoInput(self):
"""Verifies that clustering a sequential model without an input layer preserves the built state of the model."""
# No InputLayer
model = keras.Sequential([
layers.Dense(10),
layers.Dense(10),
])
self.assertEqual(model.built, False)
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(model.built, False)
# Test built state is preserved across serialization
with cluster.cluster_scope():
loaded_model = keras.models.model_from_config(
json.loads(clustered_model.to_json()))
self.assertEqual(loaded_model.built, False)
def testClusterFunctionalModelPreservesBuiltState(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)
self.assertEqual(model.built, True)
clustered_model = cluster.cluster_weights(model, **self.params)
self.assertEqual(model.built, True)
# Test built state preserves across serialization
with cluster.cluster_scope():
loaded_model = keras.models.model_from_config(
json.loads(clustered_model.to_json()))
self.assertEqual(loaded_model.built, True)
def testSameWeightsAreReturnedBeforeAndAfterSerialisation(self):
"""Verify weights of cluster_wrapper are the same after serialisation."""
# Create a dummy layer for this test
input_shape = (
1,
2,
)
original_layer = cluster_wrapper.ClusterWeights(
keras.layers.Dense(8, input_shape=input_shape),
number_of_clusters=2,
cluster_centroids_init=CentroidInitialization.LINEAR)
# Build a layer with the given shape
original_layer.build(input_shape)
model = keras.Sequential([original_layer])
# Update cluster association at least once
original_layer.update_clustered_weights_associations()
# Save and load the layer in a temp directory
with tempfile.TemporaryDirectory() as tmp_dir_name:
keras_file = os.path.join(tmp_dir_name, 'keras_model')
keras.models.save_model(model, keras_file)
with cluster.cluster_scope():
loaded_layer = keras.models.load_model(keras_file).layers[0]
def assert_list_of_variables_all_equal(l1, l2):
self.assertLen(
l1, len(l2),
'lists l1 and l2 are not equal: \n l1={l1} \n l2={l2}'.format(
l1=[v.name for v in l1], l2=[v.name for v in l2]))
name_to_var_from_l1 = {var.name: var for var in l1}
for var2 in l2:
self.assertIn(var2.name, name_to_var_from_l1)
arr1 = name_to_var_from_l1[var2.name].numpy()
arr2 = var2.numpy()
self.assertAllEqual(arr1, arr2)
# Check that trainable_weights and non_trainable_weights are the same
# in the original layer and loaded layer
assert_list_of_variables_all_equal(original_layer.trainable_weights,
loaded_layer.trainable_weights)
assert_list_of_variables_all_equal(
original_layer.non_trainable_weights,
loaded_layer.non_trainable_weights)
def test_clustered_model(clustered_model, x_test, y_test):
# Ensure accuracy persists after serializing/deserializing the model
clustered_model.save('clustered_model.h5')
# To deserialize the clustered model, use the clustering scope
with cluster.cluster_scope():
loaded_clustered_model = keras.models.load_model('clustered_model.h5')
# Checking that the deserialized model's accuracy matches the clustered model
score = loaded_clustered_model.evaluate(x_test, y_test, verbose=0)
print('Deserialized model test loss:', score[0])
print('Deserialized model test accuracy:', score[1])
# Ensure accuracy persists after stripping the model
stripped_model = cluster.strip_clustering(loaded_clustered_model)
stripped_model.compile(
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
optimizer='adam',
metrics=['accuracy'])
# Checking that the stripped model's accuracy matches the clustered model
score = stripped_model.evaluate(x_test, y_test, verbose=0)
print('Stripped model test loss:', score[0])
print('Stripped model test accuracy:', score[1])
def train_and_save(models, x_train, y_train, x_test, y_test):
for model in models:
model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'])
# Print the model summary.
model.summary()
# Model needs to be clustered after initial training
# and having achieved good accuracy
model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
print('Clustering model')
clustering_params = {
'number_of_clusters': 8,
'cluster_centroids_init': '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.categorical_crossentropy,
optimizer=opt,
metrics=['accuracy'])
# Fine-tune model
clustered_model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs_fine_tuning,
verbose=1,
validation_data=(x_test, y_test))
score = clustered_model.evaluate(x_test, y_test, verbose=0)
print('Clustered Model Test loss:', score[0])
print('Clustered Model Test accuracy:', score[1])
#Ensure accuracy persists after stripping the model
stripped_model = cluster.strip_clustering(clustered_model)
stripped_model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'])
stripped_model.save('stripped_model.h5')
# To acquire the stripped model,
# deserialize with clustering scope
with cluster.cluster_scope():
loaded_model = keras.models.load_model('stripped_model.h5')
# Checking that the stripped model's accuracy matches the clustered model
score = loaded_model.evaluate(x_test, y_test, verbose=0)
print('Stripped Model Test loss:', score[0])
print('Stripped Model Test accuracy:', score[1])
def train_and_save(models, x_train, y_train, x_test, y_test):
for model in models:
model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'])
# Print the model summary.
model.summary()
# Model needs to be clustered after initial training
# and having achieved good accuracy
model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs,
verbose=1,
validation_data=(x_test, y_test))
score = model.evaluate(x_test, y_test, verbose=0)
print('Test loss:', score[0])
print('Test accuracy:', score[1])
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.categorical_crossentropy,
optimizer=opt,
metrics=['accuracy'])
# Add callback for tensorboard summaries
log_dir = os.path.join(
FLAGS.output_dir,
datetime.datetime.now().strftime("%Y%m%d-%H%M%S-clustering"))
callbacks = [
clustering_callbacks.ClusteringSummaries(
log_dir,
cluster_update_freq='epoch',
update_freq='batch',
histogram_freq=1)
]
# Fine-tune model
clustered_model.fit(x_train,
y_train,
batch_size=batch_size,
epochs=epochs_fine_tuning,
verbose=1,
callbacks=callbacks,
validation_data=(x_test, y_test))
score = clustered_model.evaluate(x_test, y_test, verbose=0)
print('Clustered Model Test loss:', score[0])
print('Clustered Model Test accuracy:', score[1])
#Ensure accuracy persists after stripping the model
stripped_model = cluster.strip_clustering(clustered_model)
stripped_model.compile(loss=tf.keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy'])
stripped_model.save('stripped_model.h5')
# To acquire the stripped model,
# deserialize with clustering scope
with cluster.cluster_scope():
loaded_model = keras.models.load_model('stripped_model.h5')
# Checking that the stripped model's accuracy matches the clustered model
score = loaded_model.evaluate(x_test, y_test, verbose=0)
print('Stripped Model Test loss:', score[0])
print('Stripped Model Test accuracy:', score[1])
