julius/prototorch_models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FEATURE] Update pruning callback to re-add pruned prototypes

Browse Source
This commit is contained in:
Jensun Ravichandran
2021-06-04 15:56:46 +02:00
parent 42d974e08c
commit 20471bfb1c
2 changed files with 24 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
prototorch/models/callbacks.py
View File

@@ -10,12 +10,16 @@ class PruneLoserPrototypes(pl.Callback):
idle_epochs=10,
prune_quota_per_epoch=-1,
frequency=1,
replace=False,
initializer=None,
verbose=False):
self.threshold = threshold # minimum win ratio
self.idle_epochs = idle_epochs # epochs to wait before pruning
self.prune_quota_per_epoch = prune_quota_per_epoch
self.frequency = frequency
self.replace = replace
self.verbose = verbose
self.initializer = initializer
def on_epoch_end(self, trainer, pl_module):
if (trainer.current_epoch + 1) < self.idle_epochs:
@@ -24,17 +28,29 @@ class PruneLoserPrototypes(pl.Callback):
return None
ratios = pl_module.prototype_win_ratios.mean(dim=0)
to_prune = torch.arange(len(ratios))[ratios < self.threshold]
prune_labels = pl_module.prototype_labels[to_prune.tolist()]
if self.prune_quota_per_epoch > 0:
to_prune = to_prune[:self.prune_quota_per_epoch]
prune_labels = prune_labels[:self.prune_quota_per_epoch]
if len(to_prune) > 0:
if self.verbose:
print(f"\nPrototype win ratios: {ratios}")
print(f"Pruning prototypes at: {to_prune.tolist()}")
cur_num_protos = pl_module.num_prototypes
pl_module.remove_prototypes(indices=to_prune)
if self.replace:
if self.verbose:
print(f"Re-adding prototypes at: {to_prune.tolist()}")
labels, counts = torch.unique(prune_labels,
sorted=True,
return_counts=True)
distribution = dict(zip(labels.tolist(), counts.tolist()))
print(f"{distribution=}")
pl_module.add_prototypes(distribution=distribution,
initializer=self.initializer)
new_num_protos = pl_module.num_prototypes
if self.verbose:
print(f"`num_prototypes` reduced from {cur_num_protos} "
print(f"`num_prototypes` changed from {cur_num_protos} "
f"to {new_num_protos}.")
return True