| 572 | # backwards compatibility we use the buggy version by default, but you can |
| 573 | # specify legacy=False to fix it. |
| 574 | def __init__( |
| 575 | self, |
| 576 | n_e: int, |
| 577 | vq_embed_dim: int, |
| 578 | beta: float, |
| 579 | remap=None, |
| 580 | unknown_index: str = "random", |
| 581 | sane_index_shape: bool = False, |
| 582 | legacy: bool = True, |
| 583 | ): |
| 584 | super().__init__() |
| 585 | self.n_e = n_e |
| 586 | self.vq_embed_dim = vq_embed_dim |
| 587 | self.beta = beta |
| 588 | self.legacy = legacy |
| 589 | |
| 590 | self.embedding = nn.Embedding(self.n_e, self.vq_embed_dim) |
| 591 | self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) |
| 592 | |
| 593 | self.remap = remap |
| 594 | if self.remap is not None: |
| 595 | self.register_buffer("used", torch.tensor(np.load(self.remap))) |
| 596 | self.used: torch.Tensor |
| 597 | self.re_embed = self.used.shape[0] |
| 598 | self.unknown_index = unknown_index # "random" or "extra" or integer |
| 599 | if self.unknown_index == "extra": |
| 600 | self.unknown_index = self.re_embed |
| 601 | self.re_embed = self.re_embed + 1 |
| 602 | print( |
| 603 | f"Remapping {self.n_e} indices to {self.re_embed} indices. " |
| 604 | f"Using {self.unknown_index} for unknown indices." |
| 605 | ) |
| 606 | else: |
| 607 | self.re_embed = n_e |
| 608 | |
| 609 | self.sane_index_shape = sane_index_shape |
| 610 | |
| 611 | def remap_to_used(self, inds: torch.LongTensor) -> torch.LongTensor: |
| 612 | ishape = inds.shape |