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