def _get_clustered_model(self, preserve_sparsity): """Cluster the (sparse) model and return clustered_model.""" tf.random.set_seed(1) original_model = tf.keras.Sequential([ layers.Dense(5, activation='softmax', input_shape=(10, )), layers.Flatten(), ]) # Manually set sparsity in the Dense layer if preserve_sparsity is on if preserve_sparsity: first_layer_weights = original_model.layers[0].get_weights() first_layer_weights[0][:][0:2] = 0.0 original_model.layers[0].set_weights(first_layer_weights) # Start the sparsity-aware clustering clustering_params = { 'number_of_clusters': 4, 'cluster_centroids_init': cluster_config.CentroidInitialization.LINEAR, 'preserve_sparsity': True } clustered_model = experimental_cluster.cluster_weights( original_model, **clustering_params) return clustered_model
def _cluster_model(model, number_of_clusters, preserve_sparsity=False): (x_train, y_train), _ = _get_dataset() clustering_params = { 'number_of_clusters': number_of_clusters, 'cluster_centroids_init': cluster_config.CentroidInitialization.KMEANS_PLUS_PLUS, 'preserve_sparsity': preserve_sparsity, } # Cluster model clustered_model = experimental_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 testClusterKerasClusterableLayerWithSparsityPreservation(self): """Verifies that a built-in keras layer marked as clusterable is being clustered correctly when sparsity preservation is enabled.""" preserve_sparsity_params = {'preserve_sparsity': True} params = {**self.params, **preserve_sparsity_params} wrapped_layer = experimental_cluster.cluster_weights( self.keras_clusterable_layer, **params) self._validate_clustered_layer(self.keras_clusterable_layer, wrapped_layer)
def testSparsityIsPreservedDuringTraining(self): """Set a specific random seed. Ensures that we get some null weights to test sparsity preservation with. """ tf.random.set_seed(1) # Verifies that training a clustered model with null weights in it # does not destroy the sparsity of the weights. original_model = keras.Sequential([ layers.Dense(5, input_shape=(5, )), layers.Flatten(), ]) # Reset the kernel weights to reflect potential zero drifting of # the cluster centroids first_layer_weights = original_model.layers[0].get_weights() first_layer_weights[0][:][0:2] = 0.0 first_layer_weights[0][:][3] = [-0.13, -0.08, -0.05, 0.005, 0.13] first_layer_weights[0][:][4] = [-0.13, -0.08, -0.05, 0.005, 0.13] original_model.layers[0].set_weights(first_layer_weights) clustering_params = { "number_of_clusters": 6, "cluster_centroids_init": CentroidInitialization.LINEAR, "preserve_sparsity": True } clustered_model = experimental_cluster.cluster_weights( original_model, **clustering_params) stripped_model_before_tuning = cluster.strip_clustering( clustered_model) nr_of_unique_weights_before = self._get_number_of_unique_weights( stripped_model_before_tuning, 0, "kernel") clustered_model.compile( loss=keras.losses.categorical_crossentropy, optimizer="adam", metrics=["accuracy"], ) clustered_model.fit(x=self.dataset_generator(), steps_per_epoch=100) stripped_model_after_tuning = cluster.strip_clustering(clustered_model) weights_after_tuning = stripped_model_after_tuning.layers[0].kernel nr_of_unique_weights_after = self._get_number_of_unique_weights( stripped_model_after_tuning, 0, "kernel") # Check after sparsity-aware clustering, despite zero centroid can drift, # the final number of unique weights remains the same self.assertLessEqual(nr_of_unique_weights_after, nr_of_unique_weights_before) # Check that the null weights stayed the same before and after tuning. # There might be new weights that become zeros but sparsity-aware # clustering preserves the original null weights in the original positions # of the weight array self.assertTrue( np.array_equal(first_layer_weights[0][:][0:2], weights_after_tuning[:][0:2])) # Check that the number of unique weights matches the number of clusters. self.assertLessEqual(nr_of_unique_weights_after, clustering_params["number_of_clusters"])
def testClusterCustomClusterableLayerWithSparsityPreservation(self): """Verifies that a custom clusterable layer is being clustered correctly when sparsity preservation is enabled.""" preserve_sparsity_params = {'preserve_sparsity': True} params = {**self.params, **preserve_sparsity_params} wrapped_layer = experimental_cluster.cluster_weights( self.custom_clusterable_layer, **params) self.model.add(wrapped_layer) self.model.build(input_shape=(10, 1)) self._validate_clustered_layer(self.custom_clusterable_layer, wrapped_layer) self.assertEqual([('kernel', wrapped_layer.layer.kernel)], wrapped_layer.layer.get_clusterable_weights())
def testClusterModelValidLayersSuccessfulWithSparsityPreservation(self): """Verifies that clustering a sequential model results in all clusterable layers within the model being clustered when sparsity preservation is enabled.""" preserve_sparsity_params = {'preserve_sparsity': True} params = {**self.params, **preserve_sparsity_params} model = keras.Sequential([ self.keras_clusterable_layer, self.keras_non_clusterable_layer, self.custom_clusterable_layer ]) clustered_model = experimental_cluster.cluster_weights(model, **params) clustered_model.build(input_shape=(1, 28, 28, 1)) self.assertEqual(len(model.layers), len(clustered_model.layers)) for layer, clustered_layer in zip(model.layers, clustered_model.layers): self._validate_clustered_layer(layer, clustered_layer)
def testSparsityIsPreservedDuringTraining(self): """Set a specific random seed to ensure that we get some null weights to test sparsity preservation with.""" tf.random.set_seed(1) # Verifies that training a clustered model does not destroy the sparsity of # the weights. original_model = keras.Sequential([ layers.Dense(5, input_shape=(5, )), layers.Dense(5), ]) # Using a mininum number of centroids to make it more likely that some # weights will be zero. clustering_params = { "number_of_clusters": 3, "cluster_centroids_init": CentroidInitialization.LINEAR, "preserve_sparsity": True } clustered_model = experimental_cluster.cluster_weights( original_model, **clustering_params) stripped_model_before_tuning = cluster.strip_clustering( clustered_model) weights_before_tuning = stripped_model_before_tuning.get_weights()[0] non_zero_weight_indices_before_tuning = np.nonzero( weights_before_tuning) clustered_model.compile( loss=keras.losses.categorical_crossentropy, optimizer="adam", metrics=["accuracy"], ) clustered_model.fit(x=self.dataset_generator2(), steps_per_epoch=1) stripped_model_after_tuning = cluster.strip_clustering(clustered_model) weights_after_tuning = stripped_model_after_tuning.get_weights()[0] non_zero_weight_indices_after_tuning = np.nonzero(weights_after_tuning) weights_as_list_after_tuning = weights_after_tuning.reshape( -1, ).tolist() unique_weights_after_tuning = set(weights_as_list_after_tuning) # Check that the null weights stayed the same before and after tuning. self.assertTrue( np.array_equal(non_zero_weight_indices_before_tuning, non_zero_weight_indices_after_tuning)) # Check that the number of unique weights matches the number of clusters. self.assertLessEqual(len(unique_weights_after_tuning), self.params["number_of_clusters"])
def testClusterFunctionalModelSelectivelyWithSparsityPreservation(self): """Verifies that layers within a functional model can be clustered selectively when sparsity preservation is enabled.""" preserve_sparsity_params = {'preserve_sparsity': True} params = {**self.params, **preserve_sparsity_params} i1 = keras.Input(shape=(10,)) i2 = keras.Input(shape=(10,)) x1 = experimental_cluster.cluster_weights(layers.Dense(10), **params)(i1) x2 = layers.Dense(10)(i2) outputs = layers.Add()([x1, x2]) clustered_model = keras.Model(inputs=[i1, i2], outputs=outputs) self.assertIsInstance(clustered_model.layers[2], cluster_wrapper.ClusterWeights) self.assertNotIsInstance(clustered_model.layers[3], cluster_wrapper.ClusterWeights)
def testClusterSequentialModelSelectivelyWithSparsityPreservation(self): """Verifies that layers within a sequential model can be clustered selectively when sparsity preservation is enabled.""" preserve_sparsity_params = {'preserve_sparsity': True} params = {**self.params, **preserve_sparsity_params} clustered_model = keras.Sequential() clustered_model.add( experimental_cluster.cluster_weights(self.keras_clusterable_layer, **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)
def _cluster_model(original_model, sparsity_flag): """Apply the clustering wrapper, compile and train the model.""" cluster_epoch = 1 clustering_params = { 'number_of_clusters': 8, 'cluster_centroids_init': (tfmot_cluster_config.CentroidInitialization.DENSITY_BASED), 'preserve_sparsity': sparsity_flag, } cluster_model = exp_tfmot_cluster.cluster_weights(original_model, **clustering_params) callbacks = [] cluster_model = _train_model(cluster_model, callbacks, cluster_epoch) clustered_model_stripped = tfmot_cluster.strip_clustering(cluster_model) return cluster_model, clustered_model_stripped
def testClusterSimpleDenseModel(self, distribution, clustering): """End-to-end test.""" with distribution.scope(): model = experimental_cluster.cluster_weights( keras_test_utils.build_simple_dense_model(), **clustering) model.compile(loss='categorical_crossentropy', optimizer='sgd', metrics=['accuracy']) model.summary() model.fit(np.random.rand(20, 10), keras.utils.to_categorical( np.random.randint(5, size=(20, 1)), 5), epochs=1, batch_size=20) model.predict(np.random.rand(20, 10)) stripped_model = cluster.strip_clustering(model) weights_as_list = stripped_model.layers[0].kernel.numpy().reshape( -1, ).tolist() unique_weights = set(weights_as_list) self.assertLessEqual(len(unique_weights), clustering['number_of_clusters'])
def testPassingModelWithUniformWeightsToPCQAT(self, uniform_weights): """If pruned_clustered_model has uniform weights, it won't break PCQAT.""" preserve_sparsity = True original_model = tf.keras.Sequential([ layers.Dense(5, activation='softmax', input_shape=(10, )), layers.Flatten(), ]) # Manually set all weights to the same value in the Dense layer first_layer_weights = original_model.layers[0].get_weights() first_layer_weights[0][:] = uniform_weights original_model.layers[0].set_weights(first_layer_weights) # Start the sparsity-aware clustering clustering_params = { 'number_of_clusters': 4, 'cluster_centroids_init': cluster_config.CentroidInitialization.LINEAR, 'preserve_sparsity': True } clustered_model = experimental_cluster.cluster_weights( original_model, **clustering_params) clustered_model = cluster.strip_clustering(clustered_model) nr_of_unique_weights_after = self._get_number_of_unique_weights( clustered_model, 0, 'kernel') sparsity_pruning = self._get_sparsity(clustered_model) quant_aware_annotate_model = ( quantize.quantize_annotate_model(clustered_model)) sparsity_pcqat, unique_weights_pcqat = self._pcqat_training( preserve_sparsity, quant_aware_annotate_model) self.assertAllGreaterEqual(np.array(sparsity_pcqat), sparsity_pruning[0]) self.assertAllEqual(nr_of_unique_weights_after, unique_weights_pcqat)
def apply_clustering_to_conv2d(layer): if isinstance(layer, tf.keras.layers.Conv2D): return exp_tfmot_cluster.cluster_weights(layer, **clustering_params) return layer