tfm.core.actions.PruningAction

Train action to updates pruning related information.

tfm.core.actions.PruningAction(
    export_dir: str,
    model: tf.keras.Model,
    optimizer: tf.keras.optimizers.Optimizer
)

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

__call__(
    output: orbit.runner.Output
)

Update pruning step and log pruning summaries.

Args
output The train output.