def _build_clusters(self, name, layer):
"""Extracts the cluster centroids and cluster indices.
Extracts cluster centroids and cluster indices from the pretrained
clustered model when the input layer is clustered.
Args:
name: Name of weights in layer.
layer: Quantization wrapped keras layer.
Returns:
A dictionary of the initial values of the
cluster centroids, cluster indices, original weights,
the pretrained flag for marking the first training
epoch, and weight name.
"""
result = {}
weights = getattr(layer.layer, name)
if self.preserve_sparsity and not tf.reduce_any(weights == 0):
self.preserve_sparsity = False
logging.warning(
'Input layer does not contain zero weights, so apply CQAT instead.'
)
centroids_mask = None
centroids, lookup = get_unique(weights)
num_centroids = tf.size(centroids)
if self.preserve_sparsity:
sparsity_mask = tf.math.divide_no_nan(weights, weights)
zero_idx = tf.argmin(tf.abs(centroids), axis=-1)
centroids_mask = 1.0 - tf.one_hot(zero_idx, num_centroids)
result = {SPARSITY_MASK: sparsity_mask}
# Prepare clustering variables for the Keras graph when clusters
# exist, assuming we do not use number_of_clusters larger than 1024
if num_centroids > 1024:
return result
else:
clst_centroids_tf = layer.add_weight(
CLUSTER_CENTROIDS,
shape=centroids.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([centroids])[0]),
dtype=centroids.dtype,
trainable=True)
ori_weights_tf = layer.add_weight(
ORIGINAL_WEIGHTS,
shape=weights.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([weights])[0]),
dtype=weights.dtype,
trainable=True)
# Get clustering implementation according to layer type
clustering_impl_cls = clustering_registry.ClusteringLookupRegistry(
).get_clustering_impl(layer.layer, name)
clustering_impl = clustering_impl_cls(clst_centroids_tf)
pulling_indices = tf.dtypes.cast(
clustering_impl.get_pulling_indices(ori_weights_tf),
lookup.dtype)
pulling_indices_tf = layer.add_weight(
PULLING_INDICES,
shape=lookup.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([pulling_indices])[0]),
dtype=lookup.dtype,
trainable=False)
result_clst = {
CLUSTER_CENTROIDS: clst_centroids_tf,
PULLING_INDICES: pulling_indices_tf,
ORIGINAL_WEIGHTS: ori_weights_tf,
WEIGHT_NAME: name,
CLUSTERING_IMPL: clustering_impl,
CENTROIDS_MASK: centroids_mask,
}
result.update(result_clst)
return result
def build(self, input_shape):
super(ClusterWeights, self).build(input_shape)
self.build_input_shape = input_shape
# For every clusterable weights, create the clustering logic
for weight_name, weight in self.layer.get_clusterable_weights():
# Store the original weight in this wrapper
# The child reference will be overridden in
# update_clustered_weights_associations
# The actual weight_name here for the clustering wrapper is not
# necessarily the same as the original one from the layer wrapped.
# For example for cells in StackedRNNCell, the names become
# 'kernel/0', 'recurrent_kernel/0', 'kernel/1', 'recurrent_kernel/1'
original_weight = self.get_weight_from_layer(weight_name)
self.original_clusterable_weights[weight_name] = original_weight
# Track the variable
setattr(self, 'original_weight_' + weight_name, original_weight)
# Store the position in layer.weights of original_weight to restore during
# stripping
position_original_weight = next(
i for i, w in enumerate(self.layer.weights)
if w is original_weight)
self.position_original_weights[
position_original_weight] = weight_name
# Init the cluster centroids
cluster_centroids = (
clustering_centroids.CentroidsInitializerFactory.
get_centroid_initializer(self.cluster_centroids_init)(
weight, self.number_of_clusters,
self.preserve_sparsity).get_cluster_centroids())
self.cluster_centroids[weight_name] = self.add_weight(
'{}{}'.format('cluster_centroids_', weight_name),
shape=(self.number_of_clusters, ),
dtype=weight.dtype,
trainable=True,
initializer=tf.keras.initializers.Constant(
value=cluster_centroids))
# Init the weight clustering algorithm
if isinstance(self.layer, tf.keras.layers.RNN):
if isinstance(self.layer.cell,
tf.keras.layers.StackedRNNCells):
weight_name_no_index = weight_name.split('/')[0]
else:
weight_name_no_index = weight_name
elif isinstance(self.layer, tf.keras.layers.Bidirectional):
weight_name_no_index = weight_name.split('/')[0]
else:
weight_name_no_index = weight_name
self.clustering_algorithms[weight_name] = (
clustering_registry.ClusteringLookupRegistry(
).get_clustering_impl(self.layer, weight_name_no_index)(
clusters_centroids=self.cluster_centroids[weight_name],
cluster_gradient_aggregation=self.
cluster_gradient_aggregation,
))
# Init the pulling_indices (weights associations)
pulling_indices = (self.clustering_algorithms[weight_name].
get_pulling_indices(weight))
self.pulling_indices[weight_name] = self.add_weight(
'{}{}'.format('pulling_indices_', weight_name),
shape=pulling_indices.shape,
dtype=tf.int64,
trainable=False,
synchronization=tf.VariableSynchronization.ON_READ,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
initializer=tf.keras.initializers.Constant(
value=pulling_indices))
if self.preserve_sparsity:
# Init the sparsity mask
clustered_weights = (self.clustering_algorithms[weight_name].
get_clustered_weight(
pulling_indices, original_weight))
self.sparsity_masks[weight_name] = (tf.cast(tf.math.not_equal(
clustered_weights, 0),
dtype=tf.float32))
# If the model is pruned (which we suppose), this is approximately zero
self.zero_idx[weight_name] = tf.argmin(tf.abs(
self.cluster_centroids[weight_name]),
axis=-1)
def build(self, input_shape):
super(ClusterWeights, self).build(input_shape)
clusterable_weights = self.layer.get_clusterable_weights()
# Map automatically assigned TF variable name (e.g. 'dense/kernel:0') to
# provided human readable name (e.g. as in Dense(10).kernel)
clusterable_weights_to_variables = {}
for weight_name, weight in clusterable_weights:
# If a variable appears in this loop, then it is going to be removed from
# self._trainable_weights. We need to memorise what variables are going
# away so that later we are able to restore them. We have to do this to
# maintain the original order of the weights in the underlying layer.
# Incorrect order results in the incorrect OPs weights configurations.
# We can be sure that weight will be found in this array since the
# variable is either in the self._trainable_weights or in
# self._non_trainable_weights and self.weights is the result of
# concatenation of those arrays
original_index = 0
for i in range(len(self.layer.weights)):
if self.layer.weights[i].name == weight.name:
original_index = i
self.gone_variables.append(original_index)
# Again, not sure if this is needed. Leaving for now.
clusterable_weights_to_variables[self._weight_name(weight.name)] =\
weight_name
# Build initial cluster centroids for a given tensor. Factory returns a
# class and we init an object immediately
centroid_initializer = clustering_centroids.CentroidsInitializerFactory.\
get_centroid_initializer(
self.cluster_centroids_init
)(weight, self.number_of_clusters, self.preserve_sparsity)
cluster_centroids = centroid_initializer.get_cluster_centroids()
# Use k.batch_get_value since we need to initialize the variables with an
# initial value taken from a Tensor object. For each weight there is a
# different set of cluster centroids
self.cluster_centroids_tf[weight_name] = self.add_weight(
'{}{}'.format('cluster_centroids_tf_', weight_name),
shape=(self.number_of_clusters, ),
dtype=weight.dtype,
trainable=True,
initializer=initializers.Constant(
value=k.batch_get_value([cluster_centroids])[0]))
# There are vectorised implementations of look-ups, we use a new one for
# different number of dimensions.
clustering_impl_cls = clustering_registry.ClusteringLookupRegistry().\
get_clustering_impl(self.layer, weight_name)
self.clustering_impl[weight_name] = clustering_impl_cls(
self.cluster_centroids_tf[weight_name])
# We find the nearest cluster centroids and store them so that ops can
# build their weights upon it. These indices are calculated once and
# stored forever. We use to make look-ups from self.cluster_centroids_tf
pulling_indices = self.clustering_impl[weight_name].\
get_pulling_indices(weight)
self.pulling_indices_tf[weight_name] = self.add_weight(
'{}{}'.format('pulling_indices_tf_', weight_name),
shape=pulling_indices.shape,
dtype=tf.int32,
trainable=False,
synchronization=tf.VariableSynchronization.ON_READ,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
initializer=initializers.Constant(
value=k.batch_get_value([pulling_indices])[0]))
if self.preserve_sparsity:
# Get the clustered weights
clustered_weights = self.clustering_impl[
weight_name].get_clustered_weight(pulling_indices)
# Create the sparsity mask
sparsity_mask = tf.cast(tf.math.not_equal(
clustered_weights, 0),
dtype=tf.float32)
# Store the sparsity mask for training
self.sparsity_masks[weight_name] = sparsity_mask
# We store these pairs to easily update this variables later on
self.ori_weights_vars_tf[weight_name] = self.add_weight(
'{}{}'.format('ori_weights_vars_tf_', weight_name),
shape=weight.shape,
dtype=weight.dtype,
trainable=True,
initializer=initializers.Constant(
value=k.batch_get_value([weight])[0]))
# We use currying here to get an updater which can be triggered at any time
# in the future and it would return the latest version of clustered weights
def get_updater(for_weight_name):
def fn():
# Get the clustered weights
pulling_indices = self.pulling_indices_tf[for_weight_name]
clustered_weights = self.clustering_impl[for_weight_name].\
get_clustered_weight(pulling_indices)
if self.preserve_sparsity:
# Get the sparsity mask
sparsity_mask = self.sparsity_masks[for_weight_name]
# Apply the sparsity mask to the clustered weights
clustered_weights = tf.math.multiply(
clustered_weights, sparsity_mask)
return clustered_weights
return fn
# This will allow us to restore the order of weights later
# This loop stores pairs of weight names and how to restore them
for ct, weight in enumerate(self.layer.weights):
name = self._weight_name(weight.name)
full_name = '{}/{}'.format(self.layer.name, name)
if ct in self.gone_variables:
# Again, not sure if this is needed
weight_name = clusterable_weights_to_variables[name]
self.restore.append(
(name, full_name, get_updater(weight_name)))
else:
self.restore.append((name, full_name, weight))
def _build_clusters(self, name, layer):
"""Extract the cluster centroids and cluster indices from the pretrained clustered model.
Args:
name: Name of weights in layer.
layer: Quantization wrapped keras layer.
Returns:
A dictionary of the initial values of the
cluster centroids, cluster indices, original weights,
the pretrained flag for marking the first training
epoch, and weight name.
"""
weights = getattr(layer.layer, name)
centroids, lookup = get_unique(weights)
# Prepare trainable variables for the Keras graph
clst_centroids_tf = layer.add_weight(
'cluster_centroids_tf',
shape=centroids.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([centroids])[0]),
dtype=centroids.dtype,
trainable=True)
ori_weights_tf = layer.add_weight(
'ori_weights_vars_tf',
shape=weights.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([weights])[0]),
dtype=weights.dtype,
trainable=True)
# Get clustering implementation according to layer type
clustering_impl_cls = clustering_registry.ClusteringLookupRegistry().\
get_clustering_impl(layer.layer, name)
clustering_impl = clustering_impl_cls(clst_centroids_tf)
pulling_indices = tf.dtypes.cast(
clustering_impl.get_pulling_indices(ori_weights_tf),
lookup.dtype
)
pulling_indices_tf = layer.add_weight(
'pulling_indices_tf',
shape=lookup.shape,
initializer=tf.keras.initializers.Constant(
value=K.batch_get_value([pulling_indices])[0]),
dtype=lookup.dtype,
trainable=False)
for v in layer.weights:
if 'kernel' in v.name:
kernel = v
result = {
'cluster_centroids_tf': clst_centroids_tf,
'pulling_indices_tf': pulling_indices_tf,
'ori_weights_vars_tf': ori_weights_tf,
'weight_name': name,
'clst_impl': clustering_impl,
'set_kernel_weight': kernel,
}
return result
def build(self, input_shape):
super(ClusterWeights, self).build(input_shape)
clusterable_weights = self.layer.get_clusterable_weights()
# Map automatically assigned TF variable name (e.g. 'dense/kernel:0') to provided human readable name
# (e.g. as in Dense(10).kernel)
clusterable_weights_to_variables = {}
for weight_name, weight in clusterable_weights:
# If a variable appears in this loop, then it is going to be removed from self._trainable_weights.
# We need to memorise what variables are going away so that later we are able to restore them. We have to do
# this to maintain the original order of the weights in the underlying layer. Incorrect order results in the
# incorrect OPs weights configurations.
# We can be sure that weight will be found in this array since the variable is either in the
# self._trainable_weights
# or in self._non_trainable_weights and self.weights is the result of concatenation of those arrays
original_index = self.layer.weights.index(weight)
self.gone_variables.append(original_index)
# Again, not sure if this is needed. Leaving for now.
clusterable_weights_to_variables[self._weight_name(
weight.name)] = weight_name
# Build initial cluster centroids for a given tensor. Factory returns a class and we init an object immediately
centroid_initializer = clustering_centroids.CentroidsInitializerFactory.get_centroid_initializer(
self.cluster_centroids_init)(weight, self.number_of_clusters)
cluster_centroids = centroid_initializer.get_cluster_centroids()
# Use k.batch_get_value since we need to initialize the variables with an initial value taken from a Tensor object
# For each weight there is a different set of cluster centroids
self.cluster_centroids_tf[weight_name] = self.add_weight(
'cluster_centroids_tf',
shape=(self.number_of_clusters, ),
dtype=weight.dtype,
trainable=True,
initializer=initializers.Constant(
value=k.batch_get_value([cluster_centroids])[0]))
# There are vectorised implementations of look-ups, we use a new one for different number of dimensions.
clustering_impl_cls = clustering_registry.ClusteringLookupRegistry(
).get_clustering_impl(self.layer, weight_name)
self.clustering_impl[weight_name] = clustering_impl_cls(
self.cluster_centroids_tf[weight_name])
# We find the nearest cluster centroids and store them so that ops can build their weights upon it
# These indices are calculated once and stored forever. We use to make look-ups from self.cluster_centroids_tf
pulling_indices = self.clustering_impl[
weight_name].get_pulling_indices(weight)
self.pulling_indices_tf[weight_name] = self.add_weight(
'pulling_indices_tf',
shape=pulling_indices.shape,
dtype=tf.int32,
trainable=False,
initializer=initializers.Constant(
value=k.batch_get_value([pulling_indices])[0]))
# We store these pairs to easily update this variables later on
self.clustered_vars.append((weight_name, weight))
# We use currying here to get an updater which can be triggered at any time in future and it would return
# the latest version of clustered weights
def get_updater(for_weight_name):
def fn():
return self.clustering_impl[
for_weight_name].get_clustered_weight(
self.pulling_indices_tf[for_weight_name])
return fn
# This will allow us to restore the order of weights later
# This loop stores pairs of weight names and how to restore them
for ct, weight in enumerate(self.layer.weights):
name = self._weight_name(weight.name)
if ct in self.gone_variables:
# Again, not sure if this is needed
weight_name = clusterable_weights_to_variables[name]
self.restore.append((name, get_updater(weight_name)))
else:
self.restore.append((name, weight))
def build(self, input_shape):
super(ClusterWeights, self).build(input_shape)
self.build_input_shape = input_shape
# For every clusterable weights, create the clustering logic
for weight_name, weight in self.layer.get_clusterable_weights():
# Store the original weight in this wrapper
# The child reference will be overridden in
# update_clustered_weights_associations
original_weight = getattr(self.layer, weight_name)
self.original_clusterable_weights[weight_name] = original_weight
setattr(self, 'original_weight_' + weight_name,
original_weight) # Track the variable
# Store the position in layer.weights of original_weight to restore during
# stripping
position_original_weight = next(
i for i, w in enumerate(self.layer.weights)
if w is original_weight)
self.position_original_weights[
position_original_weight] = weight_name
# Init the cluster centroids
cluster_centroids = (
clustering_centroids.CentroidsInitializerFactory.
get_centroid_initializer(self.cluster_centroids_init)(
weight, self.number_of_clusters,
self.preserve_sparsity).get_cluster_centroids())
self.cluster_centroids[weight_name] = self.add_weight(
'{}{}'.format('cluster_centroids_', weight_name),
shape=(self.number_of_clusters, ),
dtype=weight.dtype,
trainable=True,
initializer=tf.keras.initializers.Constant(
value=cluster_centroids))
# Init the weight clustering algorithm
self.clustering_algorithms[weight_name] = (
clustering_registry.ClusteringLookupRegistry(
).get_clustering_impl(self.layer, weight_name)(
clusters_centroids=self.cluster_centroids[weight_name],
cluster_gradient_aggregation=self.
cluster_gradient_aggregation,
))
# Init the pulling_indices (weights associations)
pulling_indices = (self.clustering_algorithms[weight_name].
get_pulling_indices(weight))
self.pulling_indices[weight_name] = self.add_weight(
'{}{}'.format('pulling_indices_', weight_name),
shape=pulling_indices.shape,
dtype=tf.int64,
trainable=False,
synchronization=tf.VariableSynchronization.ON_READ,
aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,
initializer=tf.keras.initializers.Constant(
value=pulling_indices))
if self.preserve_sparsity:
# Init the sparsity mask
clustered_weights = (self.clustering_algorithms[weight_name].
get_clustered_weight(
pulling_indices, original_weight))
self.sparsity_masks[weight_name] = (tf.cast(tf.math.not_equal(
clustered_weights, 0),
dtype=tf.float32))
