Pruning: improve custom training loop API

[alanchiao]

The recommended path for pruning with a custom training loop is not as simple as it could be.

pruned_model = setup_pruned_model()

loss = tf.keras.losses.categorical_crossentropy
optimizer = keras.optimizers.Adam()

log_dir = tempfile.mkdtemp()

# This is all not boilerplate.
pruned_model.optimizer = optimizer
step_callback = tfmot.sparsity.keras.UpdatePruningStep()
step_callback.set_model(pruned_model)
log_callback = tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir) # optional Tensorboard logging.
log_callback.set_model(pruned_model)

step_callback.on_train_begin()
for _ in range(3):
    # only one batch given batch_size = 20 and input shape.
    step_callback.on_train_batch_begin(batch=unused_arg)
    inp = np.reshape(x_train,
                     [self._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)
    log_callback.on_epoch_end(batch=unused_arg)
...

The set_model and pruned_model.optimizer setting is unusual and could be missed.

[Naman-ntc]

I was not sure if this should be a separate issue but I found that pruning does not work unless we explicitly add training=True argument in the model call. (This seemed counterintuitive to me since from my understanding as I used to think by default training is set to be true in keras Model.)
In any case, it would likely be helpful if you specifically point out about the training argument.

[aqibsaeed]

Any updates on this to fully support custom training loop and pruning via fit function?
@alanchiao does the example code you provided work to prune model with custom training loop?

[rino20]

Hi @aqibsaeed ,

The pruning via fit function is currently available - https://www.tensorflow.org/model_optimization/guide/pruning/pruning_with_keras

I think the code above from Alan is working for any custom training loop.

[aqibsaeed]

Yes, it works thanks.

[hbellafkir]

Any updates on the custom training loop variant? This is a simple example to reproduce (tf 2.17.0):

# Define model for pruning.
pruning_params = {
    'pruning_schedule': tfmot.sparsity.keras.PolynomialDecay(initial_sparsity=0.50,
                                                             final_sparsity=0.80,
                                                             begin_step=0,
                                                             end_step=end_step)
}

model_for_pruning = prune_low_magnitude(model, **pruning_params)
update_step = tfmot.sparsity.keras.UpdatePruningStep()

for epoch in range(epochs):
    print(f'Epoch {epoch + 1}/{epochs}')

    # Iterate over the batches of the dataset
    for step, (x_batch, y_batch) in enumerate([(train_images, train_labels)]):
        with tf.GradientTape() as tape:
            # Forward pass
            logits = model_for_pruning(x_batch, training=True)
            loss_value = tf.keras.losses.sparse_categorical_crossentropy(y_batch, logits)

        # Backward pass
        grads = tape.gradient(loss_value, model_for_pruning.trainable_weights)
        model_for_pruning.optimizer.apply_gradients(zip(grads, model_for_pruning.trainable_weights))

        # Update the pruning step
        update_step.on_batch_end(step)

    update_step.on_epoch_end(batch=-1)
    val_loss, val_acc = model_for_pruning.evaluate(test_images)
    print(f'Validation loss: {val_loss}, Validation accuracy: {val_acc}')

Traceback:

Traceback (most recent call last):
  File "/workspace/xx/xy/tests/pruning.py", line 87, in <module>
    logits = model_for_pruning(x_batch, training=True)
  File "/usr/local/lib/python3.10/dist-packages/tf_keras/src/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.10/dist-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py", line 307, in call
    update_mask = utils.smart_cond(training, add_update, no_op)
  File "/usr/local/lib/python3.10/dist-packages/tensorflow_model_optimization/python/core/keras/utils.py", line 54, in smart_cond
    return smart_module.smart_cond(
  File "/usr/local/lib/python3.10/dist-packages/tensorflow_model_optimization/python/core/sparsity/keras/pruning_wrapper.py", line 290, in add_update
    tf.debugging.assert_greater_equal(
tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling layer 'prune_low_magnitude_reshape' (type PruneLowMagnitude).

Prune() wrapper requires the UpdatePruningStep callback to be provided during training. Please add it as a callback to your model.fit call.
Condition x >= y did not hold.
First 1 elements of x:
[0]
First 1 elements of y:
[1]

Call arguments received by layer 'prune_low_magnitude_reshape' (type PruneLowMagnitude):
  • inputs=tf.Tensor(shape=(60000, 28, 28), dtype=float32)
  • training=True
  • kwargs=<class 'inspect._empty'>