Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly avoids costly matrix multiplications and allows for post-hoc remapping of indices.
| 6 | from einops import rearrange |
| 7 | |
| 8 | class VectorQuantizer2(nn.Module): |
| 9 | """ |
| 10 | Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly |
| 11 | avoids costly matrix multiplications and allows for post-hoc remapping of indices. |
| 12 | """ |
| 13 | |
| 14 | # NOTE: due to a bug the beta term was applied to the wrong term. for |
| 15 | # backwards compatibility we use the buggy version by default, but you can |
| 16 | # specify legacy=False to fix it. |
| 17 | def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random", |
| 18 | sane_index_shape=False, legacy=True): |
| 19 | super().__init__() |
| 20 | self.n_e = n_e |
| 21 | self.e_dim = e_dim |
| 22 | self.beta = beta |
| 23 | self.legacy = legacy |
| 24 | |
| 25 | self.embedding = nn.Embedding(self.n_e, self.e_dim) |
| 26 | self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) |
| 27 | |
| 28 | self.remap = remap |
| 29 | if self.remap is not None: |
| 30 | self.register_buffer("used", torch.tensor(np.load(self.remap))) |
| 31 | self.re_embed = self.used.shape[0] |
| 32 | self.unknown_index = unknown_index # "random" or "extra" or integer or "closest" |
| 33 | if self.unknown_index == "extra": |
| 34 | self.unknown_index = self.re_embed |
| 35 | self.re_embed = self.re_embed + 1 |
| 36 | print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " |
| 37 | f"Using {self.unknown_index} for unknown indices.") |
| 38 | else: |
| 39 | self.re_embed = n_e |
| 40 | |
| 41 | self.sane_index_shape = sane_index_shape |
| 42 | |
| 43 | def setup_remap(self, remap, unknown_index): |
| 44 | self.remap = remap |
| 45 | self.register_buffer("used", torch.tensor(np.load(self.remap))) |
| 46 | self.re_embed = self.used.shape[0] |
| 47 | self.unknown_index = unknown_index # "random" or "extra" or integer or "closest" |
| 48 | if self.unknown_index == "extra": |
| 49 | self.unknown_index = self.re_embed |
| 50 | self.re_embed = self.re_embed + 1 |
| 51 | print(f"Remapping {self.n_e} indices to {self.re_embed} indices. " |
| 52 | f"Using {self.unknown_index} for unknown indices.") |
| 53 | |
| 54 | |
| 55 | def remap_to_used(self, inds): |
| 56 | ishape = inds.shape |
| 57 | assert len(ishape) > 1 |
| 58 | inds = inds.reshape(ishape[0], -1) |
| 59 | used = self.used.to(inds) |
| 60 | match = (inds[:, :, None] == used[None, None, ...]).long() |
| 61 | new = match.argmax(-1) |
| 62 | unknown = match.sum(2) < 1 |
| 63 | if self.unknown_index == "random": |
| 64 | new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) |
| 65 | else: |
nothing calls this directly
no outgoing calls
no test coverage detected