But this does not generalize when the num_codebooks parameters is set to 2 or more.
I think if the utilization of codebook 1 is 97% and codebook 2 is 93%, this should report (97+93)/2 = 95%, for example
Possible, ugly solution (for a single item, not a batch)
def get_codebook_util(self, indices: Tensor):
codebooks = [indices[:, :, i] for i in range(self.model.num_codebooks)]
uniques = [codebook.unique().numel() for codebook in codebooks]
mean_unique = torch.tensor(uniques, dtype=torch.float).mean()
return mean_unique / num_codes * 100
In the examples, you report
f"active %: {indices.unique().numel() / num_codes * 100:.3f}"
, for example: https://github.com/lucidrains/vector-quantize-pytorch/blob/master/examples/autoencoder_fsq.py#L76But this does not generalize when the
num_codebooks
parameters is set to 2 or more. I think if the utilization of codebook 1 is 97% and codebook 2 is 93%, this should report (97+93)/2 = 95%, for examplePossible, ugly solution (for a single item, not a batch)