Examples Codebook Utilization does not generalize

[AmitMY]

trafficstars

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#L76

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