vector_quantize_pytorch_basic_codebook_embedding_quickstart.py

1import torch
2from vector_quantize_pytorch import VectorQuantize
3
4vq = VectorQuantize(
5    dim = 256,
6    codebook_size = 512,     # codebook size
7    decay = 0.8,             # the exponential moving average decay, used if the 'commitment' metric is not used
8    commitment_weight = 1.   # the weight on the commitment loss
9)
10
11x = torch.randn(1, 1024, 256)
12quantized, indices, commit_loss = vq(x) # (1, 1024, 256), (1, 1024), (1)