tfm.core.actions.PruningAction

Train action to updates pruning related information.

This action updates pruning steps at the end of trainig loop, and log pruning metrics to tensorboard.

This action must be used when training a pruned model to avoid pruning error.

export_dir str for the export directory of the pruning summaries.
model tf.keras.Model model instance used for training. This will be used to assign a pruning step to each prunable weight.
optimizer tf.keras.optimizers.Optimizer optimizer instance used for training. This will be used to find the current training steps.

Methods

__call__

View source

Update pruning step and log pruning summaries.

Args
output The train output.