LogitEntropyLogger#

class LogitEntropyLogger(logits_key='teacher_logits', log_frequency=1, prefix='entropy', reduction='mean', log_per_class=False, default_temperature=0.07)[source]#

Bases: HookBase

Log the entropy of teacher logits using Sonata’s temperature schedule.

Monitors prediction confidence by computing the Shannon entropy of the temperature-scaled softmax over logits_key in the model output dict. In before_train it locates the Sonata submodule (unwrapping DDP) to read its current teacher_temp; if none is found it falls back to default_temperature. Runs in after_step every log_frequency steps: logs the entropy (and the temperature) to the writer under train_batch/{prefix} and to trainer.storage/iter_info, and stashes the temperature back into the model output dict for other consumers. after_epoch logs the epoch-average entropy under train/{prefix}. Registered as LogitEntropyLogger.

Parameters:

  • logits_key (str) – Key in the model output dict holding the logits. Defaults to "teacher_logits".

  • log_frequency (int) – Compute and log every this many steps. Defaults to 1.

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

  • reduction (str) – How to reduce the per-token entropy: "mean", "sum", or "none". Defaults to "mean".

  • log_per_class (bool) – If True (and reduction="none"), also log per-class mean entropy. Defaults to False.

  • default_temperature (float) – Temperature used when no Sonata model is found. Defaults to 0.07.

Note

Higher entropy indicates more uncertain predictions; lower entropy more confident ones. No-ops on a given step if the logits key is absent from the model output dict.

Example

Add to cfg.hooks for Sonata-style SSL; every log_frequency steps it logs the entropy of the teacher logits:

hooks = [dict(type="LogitEntropyLogger", log_frequency=50)]
# → writes "train_batch/entropy" and "entropy/temperature" to the writer
#   (and the epoch average "train/entropy"), puts "entropy" in
#   trainer.storage, and appends "entropy: <val>" to the iter log line
after_epoch()[source]#

Log epoch average entropy.

after_step()[source]#

Calculate and log entropy after each step.

before_train()[source]#

Initialize entropy tracking and find Sonata model.