View source on GitHub |
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__
__call__(
output: orbit.runner.Output
)
Update pruning step and log pruning summaries.
Args | |
---|---|
output
|
The train output. |