def setUp(self):
self.model = keras.Sequential()
self.params = {
'pruning_schedule': pruning_schedule.ConstantSparsity(0.5, 0),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
def setUp(self):
super(PruneTest, self).setUp()
# Layers passed in for Pruning can either be standard Keras layers provided
# by the tf.keras API (these fall under the `keras.layers` namespace), or
# custom layers provided by the user which inherit the base
# `keras.layers.Layer`.
# Standard Keras layers can either be Prunable (we know how to prune them),
# Non-Prunable (we know these layers can't be pruned) and Unsupported (we
# don't know how to deal with these yet.). Unsupported layers will raise an
# error when tried to prune.
# Custom Layers can either be Prunable (ie., they implement the
# `PrunableLayer` interface, or Non-Prunable (they don't expose any pruning
# behavior.)
# TODO(pulkitb): Change naming of Prunable/Non-Prunable/Unsupported to be
# clearer.
self.keras_prunable_layer = layers.Dense(10) # Supports pruning
self.keras_non_prunable_layer = layers.Dropout(
0.4) # Pruning not applicable
self.keras_unsupported_layer = layers.ConvLSTM2D(2,
(5, 5)) # Unsupported
self.custom_prunable_layer = CustomPrunableLayer(10)
self.custom_non_prunable_layer = CustomNonPrunableLayer(10)
self.model = keras.Sequential()
self.params = {
'pruning_schedule': pruning_schedule.ConstantSparsity(0.5, 0),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
def setUp(self):
super(PruningTest, self).setUp()
self.block_size = (1, 1)
self.block_pooling_type = "AVG"
self.constant_sparsity = pruning_schedule.ConstantSparsity(
0.5, 0, 100, 1)
def setUp(self):
super(PruneDistributedTest, self).setUp()
self.params = {
'pruning_schedule': pruning_schedule.ConstantSparsity(0.5, 0, -1, 1),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
def testExtremelySparseMask(self):
weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")
weight_dtype = weight.dtype.base_dtype
mask = tf.Variable(tf.ones(weight.get_shape(), dtype=weight_dtype),
name="mask",
dtype=weight_dtype)
threshold = tf.Variable(tf.zeros([], dtype=weight_dtype),
name="threshold",
dtype=weight_dtype)
self.initialize()
extreme_sparsity = pruning_schedule.ConstantSparsity(0.9999, 0, 100, 1)
p = pruning_impl.Pruning(pruning_vars=[(weight, mask, threshold)],
training_step_fn=self.training_step_fn,
pruning_schedule=extreme_sparsity,
block_size=self.block_size,
block_pooling_type=self.block_pooling_type)
mask_before_pruning = K.get_value(mask)
self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)
if tf.executing_eagerly():
p.conditional_mask_update()
else:
K.get_session().run(p.conditional_mask_update())
# We should always have a single connection remaining.
mask_after_pruning = K.get_value(mask)
self.assertAllEqual(np.count_nonzero(mask_after_pruning), 1)
def testPrunesModel_CustomTrainingLoop_ReachesTargetSparsity(self):
pruned_model = prune.prune_low_magnitude(
keras_test_utils.build_simple_dense_model(),
pruning_schedule=pruning_schedule.ConstantSparsity(
target_sparsity=0.5, begin_step=0, frequency=1))
batch_size = 20
x_train = np.random.rand(20, 10)
y_train = keras.utils.to_categorical(
np.random.randint(5, size=(batch_size, 1)), 5)
loss = keras.losses.categorical_crossentropy
optimizer = keras.optimizers.SGD()
unused_arg = -1
step_callback = pruning_callbacks.UpdatePruningStep()
step_callback.set_model(pruned_model)
pruned_model.optimizer = optimizer
step_callback.on_train_begin()
# 2 epochs
for _ in range(2):
step_callback.on_train_batch_begin(batch=unused_arg)
inp = np.reshape(x_train, [batch_size, 10]) # original shape: from [10].
with tf.GradientTape() as tape:
logits = pruned_model(inp, training=True)
loss_value = loss(y_train, logits)
grads = tape.gradient(loss_value, pruned_model.trainable_variables)
optimizer.apply_gradients(zip(grads, pruned_model.trainable_variables))
step_callback.on_epoch_end(batch=unused_arg)
test_utils.assert_model_sparsity(self, 0.5, pruned_model)
def setUp(self):
super(PruneIntegrationTest, self).setUp()
self.params = {
'pruning_schedule': pruning_schedule.ConstantSparsity(0.5, 0, -1, 1),
# TODO(pulkitb): Add tests for block sparsity.
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
def testSparsityValueIsValid(self, schedule_type):
if schedule_type == 'constant_sparsity':
# pylint: disable=unnecessary-lambda
self._validate_sparsity(lambda s: pruning_schedule.ConstantSparsity(s, 0))
elif schedule_type == 'polynomial_decay':
self._validate_sparsity(
lambda s: pruning_schedule.PolynomialDecay(s, 0.8, 0, 10))
self._validate_sparsity(
lambda s: pruning_schedule.PolynomialDecay(0.2, s, 0, 10))
def _construct_pruning_schedule(
schedule_type, begin_step, end_step, frequency=10):
# Uses default values for sparsity. We're only testing begin_step, end_step
# and frequency here.
if schedule_type == 'constant_sparsity':
return pruning_schedule.ConstantSparsity(
0.5, begin_step, end_step, frequency)
elif schedule_type == 'polynomial_decay':
return pruning_schedule.PolynomialDecay(
0.2, 0.8, begin_step, end_step, 3, frequency)
def testPruneWithHighSparsity_Fails(self):
params = self.params
params['pruning_schedule'] = pruning_schedule.ConstantSparsity(
target_sparsity=0.99, begin_step=0, frequency=1)
model = prune.prune_low_magnitude(
keras_test_utils.build_simple_dense_model(), **params)
with self.assertRaises(tf.errors.InvalidArgumentError):
self._train_model(model, epochs=1)
def testPrunesWithConstantSparsity(self):
sparsity = pruning_schedule.ConstantSparsity(0.5, 100, 200, 10)
step_100 = tf.Variable(100)
step_110 = tf.Variable(110)
step_200 = tf.Variable(200)
compat.initialize_variables(self)
self.assertAllClose(0.5, self.evaluate(sparsity(step_100))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_110))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_200))[1])
def testPrunesWithConstantSparsity(self):
sparsity = pruning_schedule.ConstantSparsity(0.5, 100, 200, 10)
step_100 = variables.Variable(100)
step_110 = variables.Variable(110)
step_200 = variables.Variable(200)
self.evaluate(variables.global_variables_initializer())
self.assertAllClose(0.5, self.evaluate(sparsity(step_100))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_110))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_200))[1])
def setUp(self):
super(PruningTest, self).setUp()
self.global_step = K.zeros([], dtype=dtypes.int32)
self.block_size = (1, 1)
self.block_pooling_type = "AVG"
def training_step_fn():
return self.global_step
self.constant_sparsity = pruning_schedule.ConstantSparsity(
0.5, 0, 100, 1)
self.training_step_fn = training_step_fn
def GetPrunedModel(model,
target_sparsity=0.75,
begin_step=2000,
frequency=100):
pruning_params = {
"pruning_schedule":
pruning_schedule.ConstantSparsity(target_sparsity,
begin_step=begin_step,
frequency=frequency)
}
pmodel = prune.prune_low_magnitude(model, **pruning_params)
return pmodel
def testPrunesForeverIfEndStepIsNegativeOne(self):
sparsity = pruning_schedule.ConstantSparsity(0.5, 0, -1, 10)
step_10000 = tf.Variable(10000)
step_100000000 = tf.Variable(100000000)
compat.initialize_variables(self)
self.assertTrue(self.evaluate(sparsity(step_10000))[0])
self.assertTrue(self.evaluate(sparsity(step_100000000))[0])
self.assertAllClose(0.5, self.evaluate(sparsity(step_10000))[1])
self.assertAllClose(0.5, self.evaluate(sparsity(step_100000000))[1])
def testSerializeDeserialize(self):
sparsity = pruning_schedule.ConstantSparsity(0.7, 10, 20, 10)
config = sparsity.get_config()
sparsity_deserialized = tf.keras.utils.deserialize_keras_object(
config,
custom_objects={
'ConstantSparsity': pruning_schedule.ConstantSparsity,
'PolynomialDecay': pruning_schedule.PolynomialDecay
})
self.assertEqual(sparsity.__dict__, sparsity_deserialized.__dict__)
def _prune_model(original_model):
"""Apply the pruning wrapper, compile and train the model."""
prune_epoch = 1
pruning_params = {
'pruning_schedule':
pruning_schedule.ConstantSparsity(0.50, begin_step=0, frequency=10)
}
pruning_model = prune.prune_low_magnitude(original_model, **pruning_params)
callbacks = [pruning_callbacks.UpdatePruningStep()]
pruning_model = _train_model(pruning_model, callbacks, prune_epoch)
pruning_model_stripped = prune.strip_pruning(pruning_model)
return pruning_model, pruning_model_stripped
def testPruneWithHighSparsity(self):
params = self.params
params['pruning_schedule'] = pruning_schedule.ConstantSparsity(
target_sparsity=0.99, begin_step=0, frequency=1)
model = prune.prune_low_magnitude(
keras_test_utils.build_simple_dense_model(), **params)
self._train_model(model, epochs=1)
for layer in model.layers:
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
for weight in layer.layer.get_prunable_weights():
self.assertEqual(1,
np.count_nonzero(tf.keras.backend.get_value(weight)))
def testPrunesZeroSparsity_IsNoOp(self):
model = keras_test_utils.build_simple_dense_model()
model2 = keras_test_utils.build_simple_dense_model()
model2.set_weights(model.get_weights())
params = self.params
params['pruning_schedule'] = pruning_schedule.ConstantSparsity(
target_sparsity=0, begin_step=0, frequency=1)
pruned_model = prune.prune_low_magnitude(model2, **params)
x_train = np.random.rand(20, 10),
y_train = keras.utils.to_categorical(np.random.randint(5, size=(20, 1)), 5)
self._train_model(model, epochs=1, x_train=x_train, y_train=y_train)
self._train_model(pruned_model, epochs=1, x_train=x_train, y_train=y_train)
self._assert_weights_different_objects(model, pruned_model)
self._assert_weights_same_values(model, pruned_model)
def get_pruning_params(mode='prune',
initial_sparsity=0.0,
final_sparsity=0.8,
begin_step=2000,
end_step=4000,
frequency=200):
"""Gets pruning hyper-parameters."""
p_params = {}
if mode == 'prune':
p_params['pruning_schedule'] = pruning_schedule.PolynomialDecay(
initial_sparsity=initial_sparsity,
final_sparsity=final_sparsity,
begin_step=begin_step,
end_step=end_step,
frequency=frequency)
elif mode == 'constant':
p_params['pruning_schedule'] = pruning_schedule.ConstantSparsity(
target_sparsity=final_sparsity, begin_step=begin_step)
else:
raise ValueError('Mode: %s, is not valid' % mode)
return p_params
def main():
# input image dimensions
img_rows, img_cols = 28, 28
# the data, shuffled and split between train and test sets
(x_train, y_train), (x_test, y_test) = mnist.load_data()
if K.image_data_format() == "channels_first":
x_train = x_train.reshape(x_train.shape[0], 1, img_rows, img_cols)
x_test = x_test.reshape(x_test.shape[0], 1, img_rows, img_cols)
input_shape = (1, img_rows, img_cols)
else:
x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
input_shape = (img_rows, img_cols, 1)
x_train = x_train.astype("float32")
x_test = x_test.astype("float32")
x_train /= 255
x_test /= 255
print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
# convert class vectors to binary class matrices
y_train = to_categorical(y_train, num_classes)
y_test = to_categorical(y_test, num_classes)
pruning_params = {
"pruning_schedule":
pruning_schedule.ConstantSparsity(0.75, begin_step=2000, frequency=100)
}
if prune_whole_model:
model = build_model(input_shape)
model = prune.prune_low_magnitude(model, **pruning_params)
else:
model = build_layerwise_model(input_shape, **pruning_params)
train_and_save(model, x_train, y_train, x_test, y_test)
def prune_model(original_model, train_images, train_labels):
batch_size = 256
epochs = 5
pruning_params = {
'pruning_schedule':
pruning_schedule.ConstantSparsity(0.75, begin_step=0, frequency=100)
}
pruning_model = prune.prune_low_magnitude(original_model, **pruning_params)
pruning_model.summary()
callbacks = [pruning_callbacks.UpdatePruningStep()]
fit_kwargs = {
'batch_size': batch_size,
'epochs': epochs,
'callbacks': callbacks,
}
compile_and_fit(pruning_model,
train_images,
train_labels,
compile_kwargs={},
fit_kwargs=fit_kwargs)
return pruning_model
def __init__(self,
layer,
pruning_schedule=pruning_sched.ConstantSparsity(0.5, 0),
block_size=(1, 1),
block_pooling_type='AVG',
**kwargs):
"""Create a pruning wrapper for a keras layer.
#TODO(pulkitb): Consider if begin_step should be 0 by default.
Args:
layer: The keras layer to be pruned.
pruning_schedule: A `PruningSchedule` object that controls pruning rate
throughout training.
block_size: (optional) The dimensions (height, weight) for the block
sparse pattern in rank-2 weight tensors.
block_pooling_type: (optional) The function to use to pool weights in the
block. Must be 'AVG' or 'MAX'.
**kwargs: Additional keyword arguments to be passed to the keras layer.
"""
self.pruning_schedule = pruning_schedule
self.block_size = block_size
self.block_pooling_type = block_pooling_type
# An instance of the Pruning class. This class contains the logic to prune
# the weights of this layer.
self.pruning_obj = None
# A list of all (weight,mask,threshold) tuples for this layer
self.pruning_vars = []
if block_pooling_type not in ['AVG', 'MAX']:
raise ValueError(
'Unsupported pooling type \'{}\'. Should be \'AVG\' or \'MAX\'.'
.format(block_pooling_type))
if not isinstance(layer, tf.keras.layers.Layer):
raise ValueError(
'Please initialize `Prune` layer with a '
'`Layer` instance. You passed: {input}'.format(input=layer))
# TODO(pulkitb): This should be pushed up to the wrappers.py
# Name the layer using the wrapper and underlying layer name.
# Prune(Dense) becomes prune_dense_1
kwargs.update({
'name':
'{}_{}'.format(
generic_utils.to_snake_case(self.__class__.__name__),
layer.name)
})
if isinstance(layer, prunable_layer.PrunableLayer) or hasattr(
layer, 'get_prunable_weights'):
# Custom layer in client code which supports pruning.
super(PruneLowMagnitude, self).__init__(layer, **kwargs)
elif prune_registry.PruneRegistry.supports(layer):
# Built-in keras layers which support pruning.
super(PruneLowMagnitude, self).__init__(
prune_registry.PruneRegistry.make_prunable(layer), **kwargs)
else:
raise ValueError(
'Please initialize `Prune` with a supported layer. Layers should '
'either be supported by the PruneRegistry (built-in keras layers) or '
'should be a `PrunableLayer` instance, or should has a customer '
'defined `get_prunable_weights` method. You passed: '
'{input}'.format(input=layer.__class__))
self._track_trackable(layer, name='layer')
# TODO(yunluli): Work-around to handle the first layer of Sequential model
# properly. Can remove this when it is implemented in the Wrapper base
# class.
#
# Enables end-user to prune the first layer in Sequential models, while
# passing the input shape to the original layer.
#
# tf.keras.Sequential(
# prune_low_magnitude(tf.keras.layers.Dense(2, input_shape=(3,)))
# )
#
# as opposed to
#
# tf.keras.Sequential(
# prune_low_magnitude(tf.keras.layers.Dense(2), input_shape=(3,))
# )
#
# Without this code, the pruning wrapper doesn't have an input
# shape and being the first layer, this causes the model to not be
# built. Being not built is confusing since the end-user has passed an
# input shape.
if not hasattr(self, '_batch_input_shape') and hasattr(
layer, '_batch_input_shape'):
self._batch_input_shape = self.layer._batch_input_shape
metrics.MonitorBoolGauge('prune_low_magnitude_wrapper_usage').set(
layer.__class__.__name__)
kernel_regularizer=l1(0.0001)))
model.add(QActivation(activation=quantized_relu(6), name='relu1'))
model.add(
QDense(5,
name='output',
kernel_quantizer=quantized_bits(6, 0, alpha=1),
bias_quantizer=quantized_bits(6, 0, alpha=1),
kernel_initializer='lecun_uniform',
kernel_regularizer=l1(0.0001)))
model.add(Activation(activation='softmax', name='softmax'))
from tensorflow_model_optimization.python.core.sparsity.keras import prune, pruning_callbacks, pruning_schedule
from tensorflow_model_optimization.sparsity.keras import strip_pruning
pruning_params = {
"pruning_schedule":
pruning_schedule.ConstantSparsity(0.75, begin_step=2000, frequency=100)
}
model = prune.prune_low_magnitude(model, **pruning_params)
train = not os.path.exists('model_1/KERAS_check_best_model.h5')
if train:
adam = Adam(lr=0.0001)
model.compile(optimizer=adam,
loss=['categorical_crossentropy'],
metrics=['accuracy'])
callbacks = all_callbacks(stop_patience=1000,
lr_factor=0.5,
lr_patience=10,
lr_epsilon=0.000001,
lr_cooldown=2,
lr_minimum=0.0000001,
def Q_baseline_model(size, epochs, optimizer, X_training, y_training,
X_validation, y_validation, output_name):
'''
NN Model constructor with loss and accuracy plots.
Parameters
----------
size : int
Batch size used in the training process.
epochs : int
Number of epochs the model will be trained.
optimizer : keras.optimizer
Optimizer function.
X_training : Numpy array
Training data set.
y_training : Numpy array
True labels for the training set.
X_validation : Numpy array
Validation data set.
y_validation : Numpy array
True labels for the validation set.
output_name : str
Name used for saved plots.
Returns
-------
model : qkeras.sequential
QKeras model.
w : numpy array
Array of final weights used in the model for later inference.
'''
pruning = False
# create model
name = "RMSE validation"
name2 = "RMSE training"
history = History()
model = Sequential()
model.add(
QDense(60,
input_shape=(27, ),
kernel_quantizer=quantized_bits(16, 1),
bias_quantizer=quantized_bits(16, 1),
kernel_initializer='random_normal'))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu1'))
model.add(
QDense(50,
kernel_quantizer=quantized_bits(16, 1),
bias_quantizer=quantized_bits(16, 1)))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu2'))
# model.add(Dropout(rate=0.2))
model.add(
QDense(30,
kernel_quantizer=quantized_bits(16, 1),
bias_quantizer=quantized_bits(16, 1)))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu3'))
model.add(
QDense(40,
kernel_quantizer=quantized_bits(16, 1),
bias_quantizer=quantized_bits(16, 1)))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu4'))
model.add(
QDense(15,
kernel_quantizer=quantized_bits(16, 1),
bias_quantizer=quantized_bits(16, 1)))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu5'))
# model.add(QDense(80, input_shape=(27,),kernel_quantizer=quantized_bits(16,1),bias_quantizer=quantized_bits(16,1), kernel_initializer='random_normal'))
# model.add(QActivation(activation=quantized_relu(16,1), name='relu1'))
# model.add(QDense(50,kernel_quantizer=quantized_bits(16,1),bias_quantizer=quantized_bits(16,1)))
# model.add(QActivation(activation=quantized_relu(16,1), name='relu2'))
# model.add(QDense(35,kernel_quantizer=quantized_bits(16,1),bias_quantizer=quantized_bits(16,1)))
# model.add(QActivation(activation=quantized_relu(16,1), name='relu3'))
# # # model.add(Dropout(rate=0.2))
model.add(QDense(1, kernel_quantizer=quantized_bits(16, 1)))
model.add(QActivation(activation=quantized_relu(16, 1), name='relu6'))
#model.add(Activation("sigmoid"))
# model.add(QActivation(activation=quantized_tanh(16,1),name='tanh'))
if pruning == True:
print("////////////////////////Training Model with pruning")
pruning_params = {
"pruning_schedule":
pruning_schedule.ConstantSparsity(0.75,
begin_step=2000,
frequency=100)
}
model = prune.prune_low_magnitude(model, **pruning_params)
model.compile(loss='mean_squared_error', optimizer=optimizer)
model.fit(X_training,
y_training,
batch_size=size,
epochs=epochs,
verbose=1,
validation_data=(X_validation, y_validation),
callbacks=[history,
pruning_callbacks.UpdatePruningStep()])
model = strip_pruning(model)
w = model.layers[0].weights[0].numpy()
h, b = np.histogram(w, bins=100)
plt.figure(figsize=(7, 7))
plt.bar(b[:-1], h, width=b[1] - b[0])
plt.semilogy()
plt.savefig("Zeros' distribution", format='png')
print('% of zeros = {}'.format(np.sum(w == 0) / np.size(w)))
else:
print("////////////////////////Training Model WITHOUT pruning")
model.compile(loss='mean_squared_error', optimizer=optimizer)
model.fit(X_training,
y_training,
batch_size=size,
epochs=epochs,
verbose=1,
validation_data=(X_validation, y_validation),
callbacks=[history])
# Compile model
# model.compile(loss='mean_squared_error', optimizer=optimizer)
# model.fit(X_training, y_training,
# batch_size=size,
# epochs=epochs,
# verbose=1,
# validation_data=(X_validation, y_validation),callbacks=[history])
w = []
for layer in model.layers:
print(layer)
w.append(layer.get_weights())
#print(w)
train_predictions = model.predict(X_training)
predictions = model.predict(X_validation)
lin_mse = mean_squared_error(y_validation, predictions)
lin_rmse = np.sqrt(lin_mse)
lin_mse2 = mean_squared_error(y_training, train_predictions)
lin_rmse2 = np.sqrt(lin_mse2)
msg = "%s: %f" % (name, lin_rmse)
msg2 = "%s: %f" % (name2, lin_rmse2)
print(msg)
print(msg2)
fig, ax = plt.subplots()
# xy=np.vstack([y_validation, predictions])
#z=gaussian_kde(xy)
ax.scatter(y_validation, predictions, edgecolors=(0, 0, 0))
ax.set_title('Regression model predictions (validation set)')
ax.set_xlabel('Measured $p_T$ (GeV/c)')
ax.set_ylabel('Predicted $p_T$ (GeV/c)')
ax.plot([Y.min(), Y.max()], [Y.min(), Y.max()], 'k--', lw=4)
plt.rc('font', size=20)
plt.rc('axes', titlesize=18)
plt.rc('axes', labelsize=18)
plt.rc('xtick', labelsize=18)
plt.rc('ytick', labelsize=18)
plt.rc('legend', fontsize=18)
plt.rc('figure', titlesize=18)
plt.tight_layout()
plt.savefig(outrootname + '/' + '1' + output_name, format='png', dpi=800)
fig2, ax2 = plt.subplots()
ax2.plot(history.history['loss'], label='loss')
ax2.plot(history.history['val_loss'], label='val_loss')
ax2.set_title('Training and Validation loss per epoch')
ax2.set_xlabel('# Epoch')
ax2.set_ylabel('loss')
plt.legend()
plt.tight_layout()
plt.savefig(outrootname + '/' + '2' + output_name, format='png', dpi=800)
#plt.show()
del ax, ax2
return model, w
def __init__(self,
layer,
pruning_schedule=pruning_sched.ConstantSparsity(0.5, 0),
block_size=(1, 1),
block_pooling_type='AVG',
**kwargs):
"""Create a pruning wrapper for a keras layer.
#TODO(pulkitb): Consider if begin_step should be 0 by default.
Args:
layer: The keras layer to be pruned.
pruning_schedule: A `PruningSchedule` object that controls pruning rate
throughout training.
block_size: (optional) The dimensions (height, weight) for the block
sparse pattern in rank-2 weight tensors.
block_pooling_type: (optional) The function to use to pool weights in the
block. Must be 'AVG' or 'MAX'.
**kwargs: Additional keyword arguments to be passed to the keras layer.
"""
self.pruning_schedule = pruning_schedule
self.block_size = block_size
self.block_pooling_type = block_pooling_type
# An instance of the Pruning class. This class contains the logic to prune
# the weights of this layer.
self.pruning_obj = None
# A list of all (weight,mask,threshold) tuples for this layer
self.pruning_vars = []
if block_pooling_type not in ['AVG', 'MAX']:
raise ValueError(
'Unsupported pooling type \'{}\'. Should be \'AVG\' or \'MAX\'.'
.format(block_pooling_type))
if not isinstance(layer, Layer):
raise ValueError(
'Please initialize `Prune` layer with a '
'`Layer` instance. You passed: {input}'.format(input=layer))
# TODO(pulkitb): This should be pushed up to the wrappers.py
# Name the layer using the wrapper and underlying layer name.
# Prune(Dense) becomes prune_dense_1
kwargs.update({
'name':
'{}_{}'.format(
generic_utils.to_snake_case(self.__class__.__name__),
layer.name)
})
if isinstance(layer, prunable_layer.PrunableLayer):
# Custom layer in client code which supports pruning.
super(PruneLowMagnitude, self).__init__(layer, **kwargs)
elif prune_registry.PruneRegistry.supports(layer):
# Built-in keras layers which support pruning.
super(PruneLowMagnitude, self).__init__(
prune_registry.PruneRegistry.make_prunable(layer), **kwargs)
else:
raise ValueError(
'Please initialize `Prune` with a supported layer. Layers should '
'either be a `PrunableLayer` instance, or should be supported by the '
'PruneRegistry. You passed: {input}'.format(
input=layer.__class__))
self._track_trackable(layer, name='layer')
# TODO(yunluli): Work-around to handle the first layer of Sequential model
# properly. Can remove this when it is implemented in the Wrapper base
# class.
# The _batch_input_shape attribute in the first layer makes a Sequential
# model to be built. This change makes sure that when we apply the wrapper
# to the whole model, this attribute is pulled into the wrapper to preserve
# the 'built' state of the model.
if not hasattr(self, '_batch_input_shape') and hasattr(
layer, '_batch_input_shape'):
self._batch_input_shape = self.layer._batch_input_shape
def prune_low_magnitude(to_prune,
pruning_schedule=pruning_sched.ConstantSparsity(
0.5, 0),
block_size=(1, 1),
block_pooling_type='AVG',
pruning_policy=None,
**kwargs):
"""Modify a tf.keras layer or model to be pruned during training.
This function wraps a tf.keras model or layer with pruning functionality which
sparsifies the layer's weights during training. For example, using this with
50% sparsity will ensure that 50% of the layer's weights are zero.
The function accepts either a single keras layer
(subclass of `tf.keras.layers.Layer`), list of keras layers or a Sequential
or Functional tf.keras model and handles them appropriately.
If it encounters a layer it does not know how to handle, it will throw an
error. While pruning an entire model, even a single unknown layer would lead
to an error.
Prune a model:
```python
pruning_params = {
'pruning_schedule': ConstantSparsity(0.5, 0),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
model = prune_low_magnitude(
keras.Sequential([
layers.Dense(10, activation='relu', input_shape=(100,)),
layers.Dense(2, activation='sigmoid')
]), **pruning_params)
```
Prune a layer:
```python
pruning_params = {
'pruning_schedule': PolynomialDecay(initial_sparsity=0.2,
final_sparsity=0.8, begin_step=1000, end_step=2000),
'block_size': (2, 3),
'block_pooling_type': 'MAX'
}
model = keras.Sequential([
layers.Dense(10, activation='relu', input_shape=(100,)),
prune_low_magnitude(layers.Dense(2, activation='tanh'), **pruning_params)
])
```
Pretrained models: you must first load the weights and then apply the
prune API:
```python
model.load_weights(...)
model = prune_low_magnitude(model)
```
Optimizer: this function removes the optimizer. The user is expected to
compile the model
again. It's easiest to rely on the default (step starts at 0) and then
use that to determine the desired begin_step for the pruning_schedules.
Checkpointing: checkpointing should include the optimizer, not just the
weights. Pruning supports
checkpointing though
upon inspection, the weights of checkpoints are not sparse
(https://github.com/tensorflow/model-optimization/issues/206).
Arguments:
to_prune: A single keras layer, list of keras layers, or a
`tf.keras.Model` instance.
pruning_schedule: A `PruningSchedule` object that controls pruning rate
throughout training.
block_size: (optional) The dimensions (height, weight) for the block
sparse pattern in rank-2 weight tensors.
block_pooling_type: (optional) The function to use to pool weights in the
block. Must be 'AVG' or 'MAX'.
pruning_policy: (optional) The object that controls to which layers
`PruneLowMagnitude` wrapper will be applied. This API is experimental
and is subject to change.
**kwargs: Additional keyword arguments to be passed to the keras layer.
Ignored when to_prune is not a keras layer.
Returns:
Layer or model modified with pruning wrappers. Optimizer is removed.
Raises:
ValueError: if the keras layer is unsupported, or the keras model contains
an unsupported layer.
"""
def _prune_list(layers, **params):
wrapped_layers = []
for layer in layers:
# Allow layer that is already wrapped by the pruning wrapper
# to be used as is.
# No need to wrap the input layer either.
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
wrapped_layers.append(layer)
elif isinstance(layer, keras.layers.InputLayer):
# TODO(yunluli): Replace with a clone function in keras.
wrapped_layers.append(
layer.__class__.from_config(layer.get_config()))
else:
wrapped_layers.append(
pruning_wrapper.PruneLowMagnitude(layer, **params))
return wrapped_layers
def _add_pruning_wrapper(layer):
if isinstance(layer, keras.Model):
# Check whether the model is a subclass model.
if (not layer._is_graph_network
and not isinstance(layer, keras.models.Sequential)):
raise ValueError(
'Subclassed models are not supported currently.')
return keras.models.clone_model(
layer, input_tensors=None, clone_function=_add_pruning_wrapper)
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
return layer
if pruning_policy and not pruning_policy.allow_pruning(layer):
return layer
else:
return pruning_wrapper.PruneLowMagnitude(layer, **params)
params = {
'pruning_schedule': pruning_schedule,
'block_size': block_size,
'block_pooling_type': block_pooling_type
}
is_sequential_or_functional = isinstance(
to_prune, keras.Model) and (isinstance(to_prune, keras.Sequential)
or to_prune._is_graph_network)
# A subclassed model is also a subclass of keras.layers.Layer.
is_keras_layer = isinstance(
to_prune,
keras.layers.Layer) and not isinstance(to_prune, keras.Model)
if isinstance(to_prune, list):
return _prune_list(to_prune, **params)
elif is_sequential_or_functional:
if pruning_policy:
pruning_policy.ensure_model_supports_pruning(to_prune)
return _add_pruning_wrapper(to_prune)
elif is_keras_layer:
params.update(kwargs)
return pruning_wrapper.PruneLowMagnitude(to_prune, **params)
else:
raise ValueError(
'`prune_low_magnitude` can only prune an object of the following '
'types: tf.keras.models.Sequential, tf.keras functional model, '
'tf.keras.layers.Layer, list of tf.keras.layers.Layer. You passed '
'an object of type: {input}.'.format(
input=to_prune.__class__.__name__))
def prune_low_magnitude(to_prune,
pruning_schedule=pruning_sched.ConstantSparsity(
0.5, 0),
block_size=(1, 1),
block_pooling_type='AVG',
**kwargs):
"""Modify a keras layer or model to be pruned during training.
This function wraps a tf.keras model or layer with pruning functionality which
sparsifies the layer's weights during training. For example, using this with
50% sparsity will ensure that 50% of the layer's weights are zero.
The function accepts either a single keras layer
(subclass of `tf.keras.layers.Layer`), list of keras layers or a Sequential
or Functional keras model and handles them appropriately.
If it encounters a layer it does not know how to handle, it will throw an
error. While pruning an entire model, even a single unknown layer would lead
to an error.
Prune a model:
```python
pruning_params = {
'pruning_schedule': ConstantSparsity(0.5, 0),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
model = prune_low_magnitude(
keras.Sequential([
layers.Dense(10, activation='relu', input_shape=(100,)),
layers.Dense(2, activation='sigmoid')
]), **pruning_params)
```
Prune a layer:
```python
pruning_params = {
'pruning_schedule': PolynomialDecay(initial_sparsity=0.2,
final_sparsity=0.8, begin_step=1000, end_step=2000),
'block_size': (2, 3),
'block_pooling_type': 'MAX'
}
model = keras.Sequential([
layers.Dense(10, activation='relu', input_shape=(100,)),
prune_low_magnitude(layers.Dense(2, activation='tanh'), **pruning_params)
])
```
Arguments:
to_prune: A single keras layer, list of keras layers, or a
`tf.keras.Model` instance.
pruning_schedule: A `PruningSchedule` object that controls pruning rate
throughout training.
block_size: (optional) The dimensions (height, weight) for the block
sparse pattern in rank-2 weight tensors.
block_pooling_type: (optional) The function to use to pool weights in the
block. Must be 'AVG' or 'MAX'.
**kwargs: Additional keyword arguments to be passed to the keras layer.
Ignored when to_prune is not a keras layer.
Returns:
Layer or model modified with pruning wrappers.
Raises:
ValueError: if the keras layer is unsupported, or the keras model contains
an unsupported layer.
"""
def _prune_list(layers, **params):
wrapped_layers = []
for layer in layers:
# Allow layer that is already wrapped by the pruning wrapper
# to be used as is.
# No need to wrap the input layer either.
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
wrapped_layers.append(layer)
elif isinstance(layer, InputLayer):
# TODO(yunluli): Replace with a clone function in keras.
wrapped_layers.append(
layer.__class__.from_config(layer.get_config()))
else:
wrapped_layers.append(
pruning_wrapper.PruneLowMagnitude(layer, **params))
return wrapped_layers
def _add_pruning_wrapper(layer):
if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
return layer
return pruning_wrapper.PruneLowMagnitude(layer, **params)
params = {
'pruning_schedule': pruning_schedule,
'block_size': block_size,
'block_pooling_type': block_pooling_type
}
is_sequential_or_functional = isinstance(
to_prune, keras.Model) and (isinstance(to_prune, keras.Sequential)
or to_prune._is_graph_network)
# A subclassed model is also a subclass of keras.layers.Layer.
is_keras_layer = isinstance(
to_prune,
keras.layers.Layer) and not isinstance(to_prune, keras.Model)
if isinstance(to_prune, list):
return _prune_list(to_prune, **params)
elif is_sequential_or_functional:
return keras.models.clone_model(to_prune,
input_tensors=None,
clone_function=_add_pruning_wrapper)
elif is_keras_layer:
params.update(kwargs)
return pruning_wrapper.PruneLowMagnitude(to_prune, **params)
else:
raise ValueError(
'`prune_low_magnitude` can only prune an object of the following '
'types: tf.keras.models.Sequential, tf.keras functional model, '
'tf.keras.layers.Layer, list of tf.keras.layers.Layer. You passed '
'an object of type: {input}.'.format(
input=to_prune.__class__.__name__))
def apply_pruning_to_conv2D(layer):
if isinstance(layer, tf.keras.layers.Conv2D):
return tfmot.sparsity.keras.prune_low_magnitude(layer,pruning_schedule=pruning_sched.ConstantSparsity(0.8,0))
return layer
projection_simplex_sort(np.array([1, 0.1]), 3)
#%%pixie_debugger
# l1 parameters
l1coeff = 0
# for reconstructor
#l2coeff = 0
# for keras sparsity
sparsity = 0.3
pruning_params = {
'pruning_schedule': pruning_sched.ConstantSparsity(0.9, 0),
#'pruning_schedule': pruning_sched.PolynomialDecay(0, 0.3, 0, 100),
'block_size': (1, 1),
'block_pooling_type': 'AVG'
}
def component_diff_normalized(v):
"""How much the vector is close to (1,0) or (0,1)."""
v = np.abs(v)
v = np.sort(v)[::-1]
maxv = v[0]
smaxv = v[1]
return 1. - (maxv - smaxv) / maxv
