Pruning: Checkpointed Models Not Guaranteed To Be Sparse

Models on final export are sparse and pruning generally handles checkpointing correctly.

However, upon inspection, the weights of a checkpoint are not sparse themselves.

Reproduce

See testPruneCheckpoints_CheckpointsNotSparse in prune_integration_test.py

Can search for "model-optimization/issues/206" in codebase to find unit test also.

Theory

We sparsify the weights on two occasions: the UpdatePruningStep callback on epoch end and the call to self.layer.add_update(self.pruning_obj.weight_mask_op()) under the Pruning Wrapper. The default ModelCheckpoint callback checkpoints on epoch end.

The following order of events would cause the checkpointed models to not be sparse.

PruningWrapper update (weights are sparse) -> backwards propagation (weights aren't sparse) -> Checkpoint Callback -> UpdatePruningStep Callback (weights are sparse again)

This order seems feasible because we don't know of any guarantees that callbacks execute in a particular order. If both callbacks happen at the same time, then the checkpointed model could be partially sparse (e.g. weights are 30% sparse even when the mask is 50% sparse). If the checkpoint callback happens prior, then the weights wouldn't be sparse at all.

Supporting Evidence

1. When we save the model at the end after pruning, it has always compressed.
2. Previously when checkpointing was done on epoch_begin, the checkpointed models were compressed. The theory is that on_epoch_begin forces checkpointing to happen after UpdatePruningStep as follows:

PruningWrapper update (weights are sparse) -> backwards propagation (weights aren't sparse) -> UpdatePruningStep Callback (weights are sparse again) -> Checkpoint Callback

Notes

Any changes made will have to consider backwards-compatibility (e.g. existing training behavior is unchanged except
for fixing this pathway). Otherwise, it'll have to go into the next major release.