peal.adaptors.didae.ProbeModel

class peal.adaptors.didae.ProbeModel[source]

Bases: Module

Linear probe on the generator’s z_sem posing as a binary classifier.

Wraps the distilled probe (w, bias) of step 1 so it can be evaluated with calculate_test_accuracy like the student: an image in classifier normalisation is converted to generator normalisation, encoded to z_sem and scored as s = z @ w + bias; the output logits are [-s, s] so argmax gives class 1 when s > 0.

Parameters:

  • generator (InvertibleGenerator) – Generator with encode(x, only_semantic=True).

  • w (torch.Tensor) – Probe weight of shape (encoder_dim,).

  • bias (float) – Probe bias.

  • classifier_to_generator (callable) – Maps classifier-normalised images to generator normalisation.

  • device (str or torch.device) – Device the encoding runs on.

Methods

forward

Return (B, 2) logits [-s, s] for a batch of classifier inputs.
__init__(generator, w, bias, classifier_to_generator, device)[source]

Store the probe and the encoder it reads from.

See the class docstring for the meaning of the arguments.

forward(x)[source]

Return (B, 2) logits [-s, s] for a batch of classifier inputs.