PrototypeUsageLogger#

class PrototypeUsageLogger(log_frequency=10, prefix='prototypes', log_pooling_levels=True)[source]#

Bases: HookBase

Monitor prototype utilization in Sonata-style clustering heads.

In before_train (after unwrapping DDP) it finds the Sonata submodule and registers a forward hook on each teacher/student head whose name contains "head". Each forward hook fires every log_frequency calls (training mode only) and, from the head’s output logits, computes prototype-usage statistics via argmax assignment counts (synchronized across ranks): number of used prototypes, percent unused, average tokens per active prototype, and the assignment-distribution entropy. These are written to the writer under {prefix}/<head>/.... after_train removes all registered hooks. Registered as PrototypeUsageLogger.

When log_pooling_levels is set it ALSO hooks each teacher/student backbone and, by walking the pooling_parent chain of the returned Point (read-only, so the model’s later up-cast is untouched), logs the cross-rank average number of points per event at each hierarchy stage under {prefix}/<backbone>/points_per_event/level{k} (level0 = finest input resolution, increasing with pooling depth). This surfaces redundant pooling stages – e.g. a deepest level that collapses to ~1 point per event after many pooling operations.

Parameters:

  • log_frequency (int) – Log statistics every this many head/backbone forward passes. Defaults to 10.

  • prefix (str) – Namespace prefix for the logged keys. Defaults to "prototypes".

  • log_pooling_levels (bool) – Also log average points-per-event at each backbone pooling stage. Defaults to True.

Note

Specific to models exposing a Sonata-like teacher/student ModuleDict of heads; warns and does nothing if none is found. Counts are all-reduced so the reported usage is global across GPUs.

Example

Add to cfg.hooks for Sonata-style SSL; it watches each clustering head’s prototype assignments:

hooks = [dict(type="PrototypeUsageLogger", log_frequency=20)]
# → every 20 head forward passes writes, per head, the cross-rank
#   "prototypes/<teacher|student>/<head>/{used_count,unused_percent,
#   tokens_per_prototype,assignment_entropy}" to the writer, plus
#   "prototypes/<teacher|student>/backbone/points_per_event/level{k}"
#   for each pooling stage
after_train()[source]#

Clean up hooks when training is done.

before_train()[source]#

Register hooks on Sonata heads to capture prototype usage.