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.

Args
`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__(
    output: orbit.runner.Output
)

Update pruning step and log pruning summaries.

Args
`output`	The train output.