peal.sparse_dictionaries.activation_cache.extract

peal.sparse_dictionaries.activation_cache.extract(datasets, feature_extractor, device, batch_size=256, num_workers=4)[source]

Encode every dataset, returning one (X, Y) pair per dataset.

Parameters:

  • datasets (list of torch.utils.data.Dataset) – Datasets to encode, each yielding x or (x, y, ...) items.

  • feature_extractor (callable) – Maps a batch of inputs on device to activations [B, act_size].

  • device (str or torch.device) – Device the inputs are moved to before encoding.

  • batch_size (int) – DataLoader batch size.

  • num_workers (int) – DataLoader worker processes.

Returns:

[(X, Y), ...] with X a float CPU tensor [N, act_size] and Y the concatenated labels, or None when the dataset yields no label. Progress is printed every 100 batches.

Return type:

list of tuple