FeatureStdMonitor#

class FeatureStdMonitor(log_frequency=10, prefix='feature_std', monitor_student=True, monitor_teacher=True, track_channels=False)[source]#

Bases: HookBase

Monitor feature standard deviation to detect representation collapse.

In before_train (after unwrapping DDP) it locates the student/teacher module (Sonata, JEPA, etc.) and registers forward hooks on the selected backbones (and their representation_fusion if present). Each forward hook fires every log_frequency calls (training mode only) and computes, with cross-rank synchronization, the global feature std, per-sample batch std, and per-channel std summary (mean/min/max), logging them to the writer under {prefix}/<module>/...; with track_channels=True it also logs each channel’s std. after_train removes all hooks. Useful for catching feature collapse (std trending to zero). Registered as FeatureStdMonitor.

Parameters:

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

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

  • monitor_student (bool) – Register hooks on the student branch. Defaults to True.

  • monitor_teacher (bool) – Register hooks on the teacher branch. Defaults to True.

  • track_channels (bool) – If True, additionally log per-channel std. Defaults to False.

Note

Statistics are computed inside the forward pass to avoid retaining large feature tensors. Warns and does nothing if no student/teacher module is found.

Example

Add to cfg.hooks to catch representation collapse during SSL; it watches the std of teacher/student backbone features:

hooks = [dict(type="FeatureStdMonitor", log_frequency=20)]
# → every 20 forward passes writes "feature_std/<student|teacher>/
#   {global_std,batch_std,channel_mean_std,channel_min_std,
#   channel_max_std}" to the writer (std trending to 0 signals collapse)
after_train()[source]#

Clean up hooks.

before_train()[source]#

Register forward hooks to capture feature statistics.