def on_stage_end(self, runner: "IRunner") -> None:
"""
On stage end action.
Active if prune_on_stage_end or
remove_reparametrization is True.
Args:
runner: runner for your experiment
"""
if self.prune_on_stage_end:
prune_model(
model=runner.model,
pruning_fn=self.pruning_fn,
keys_to_prune=self.keys_to_prune,
amount=self.amount,
layers_to_prune=self.layers_to_prune,
reinitialize_after_pruning=self.reinitialize_after_pruning,
)
if self.remove_reparametrization_on_stage_end:
remove_reparametrization(
model=runner.model,
keys_to_prune=self.keys_to_prune,
layers_to_prune=self.layers_to_prune,
)
def on_epoch_end(self, runner: "IRunner") -> None:
"""Event handler.
Active if prune_on_epoch_end is True.
Args:
runner: runner for your experiment
"""
if self.prune_on_epoch_end and runner.stage_epoch_step != runner.stage_epoch_len:
prune_model(
model=runner.model,
pruning_fn=self.pruning_fn,
keys_to_prune=self.keys_to_prune,
amount=self.amount,
layers_to_prune=self.layers_to_prune,
)
def on_epoch_end(self, runner: "IRunner") -> None:
"""
On epoch end action.
Active if prune_on_epoch_end is True.
Args:
runner: runner for your experiment
"""
if self.prune_on_epoch_end and runner.num_epochs != runner.epoch:
prune_model(
model=runner.model,
pruning_fn=self.pruning_fn,
keys_to_prune=self.keys_to_prune,
amount=self.amount,
layers_to_prune=self.layers_to_prune,
reinitialize_after_pruning=self.reinitialize_after_pruning,
)