braindecode.modules.EMACodebook

class braindecode.modules.EMACodebook(dim, codebook_size, decay=0.99, epsilon=1e-06, threshold_ema_dead_code=2, kmeans_init=True, kmeans_iters=10)

EMA-updated Euclidean codebook.

Buffer keys match the upstream BrainOmni checkpoint for state-dict parity.

Parameters:

• dim (int) – Dimensionality of each codebook vector.

• codebook_size (int) – Number of codebook entries.

• decay (float) – Exponential moving-average decay for codebook updates.

• epsilon (float) – Small constant for numerical stability in EMA normalisation.

• threshold_ema_dead_code (int) – Minimum cluster size below which a code is considered dead and replaced with a random sample from the current batch.

• kmeans_init (bool) – If True, initialize embeddings and cluster counts from the first input batch, matching the released BrainOmni training implementation.

• kmeans_iters (int) – Number of K-means iterations used for first-batch initialization.

Methods

dequantize(embed_ind)

Return type:

Tensor

expire_codes_(batch_samples)

Return type:

None

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

tuple[Tensor, Tensor]

init_embed_(data)

Initialize fresh codebooks from data and synchronize rank-zero state.

Return type:

None

quantize(x)

Return type:

Tensor

replace_(samples, mask)

Return type:

None