Spaces:
Running on Zero
Running on Zero
File size: 5,897 Bytes
247228a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce
def log(t, eps=1e-5):
return t.clamp(min=eps).log()
def entropy(prob):
return (-prob * log(prob)).sum(dim=-1)
class VectorQuantizer(nn.Module):
def __init__(
self,
n_e,
e_dim,
beta,
normalize_embedding,
remap=None,
unknown_index="random",
sane_index_shape=False,
legacy=True,
diversity_gamma=1.0,
frac_per_sample_entropy=1.0,
token_noise=0.0,
):
super().__init__()
self.n_e = n_e
self.e_dim = e_dim
self.beta = beta
self.legacy = legacy
self.normalize_embedding = normalize_embedding
self.diversity_gamma = diversity_gamma
self.frac_per_sample_entropy = frac_per_sample_entropy
self.token_noise = token_noise
self.sane_index_shape = sane_index_shape
# Codebook
self.embedding = nn.Embedding(n_e, e_dim)
self.embedding.weight.data.uniform_(-1.0 / n_e, 1.0 / n_e)
if self.normalize_embedding:
self.embedding.weight.data = F.normalize(self.embedding.weight.data, dim=1)
# Optional remapping
self.remap = remap
if remap is not None:
self.register_buffer("used", torch.tensor(np.load(remap)))
self.re_embed = self.used.shape[0]
self.unknown_index = unknown_index
if unknown_index == "extra":
self.unknown_index = self.re_embed
self.re_embed += 1
print(
f"Remapping {n_e} indices to {self.re_embed} indices. "
f"Using {self.unknown_index} for unknown indices."
)
else:
self.re_embed = n_e
def remap_to_used(self, indices):
ishape = indices.shape
indices = indices.view(ishape[0], -1)
used = self.used.to(indices)
match = (indices[:, :, None] == used[None, None, :]).long()
new = match.argmax(-1)
unknown = match.sum(2) < 1
if self.unknown_index == "random":
new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(indices.device)
else:
new[unknown] = self.unknown_index
return new.view(ishape)
def unmap_to_all(self, indices):
ishape = indices.shape
indices = indices.view(ishape[0], -1)
used = self.used.to(indices)
if self.re_embed > used.shape[0]:
indices[indices >= used.shape[0]] = 0
gathered = torch.gather(used.expand(indices.shape[0], -1), 1, indices)
return gathered.view(ishape)
def entropy_loss(self, distances, inv_temperature=100.0):
prob = (-distances * inv_temperature).softmax(dim=-1)
if self.frac_per_sample_entropy < 1.0:
num_tokens = prob.shape[0]
sample_size = int(num_tokens * self.frac_per_sample_entropy)
mask = torch.randperm(num_tokens, device=prob.device)[:sample_size]
per_sample_probs = prob[mask]
else:
per_sample_probs = prob
per_sample_entropy = entropy(per_sample_probs).mean()
avg_prob = reduce(per_sample_probs, "... d -> d", "mean")
codebook_entropy = entropy(avg_prob).mean()
return per_sample_entropy - self.diversity_gamma * codebook_entropy
def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
assert temp in (None, 1.0)
assert not rescale_logits and not return_logits
if self.normalize_embedding:
self.embedding.weight.data = F.normalize(self.embedding.weight.data, dim=1)
# Flatten input
z = rearrange(z, "b c h w -> b h w c").contiguous()
z_flat = z.view(-1, self.e_dim)
# Compute distances
e = self.embedding.weight
d = (
torch.sum(z_flat ** 2, dim=1, keepdim=True)
+ torch.sum(e ** 2, dim=1)
- 2 * torch.einsum("bd,dn->bn", z_flat, e.T)
)
min_indices = torch.argmin(d, dim=1)
# Optional token noise
if self.token_noise > 0.0 and self.training:
noise_mask = torch.rand_like(min_indices.float()) < self.token_noise
rand_indices = torch.randint(0, self.n_e, min_indices.shape, device=z.device)
min_indices[noise_mask] = rand_indices[noise_mask]
z_q = self.embedding(min_indices).view_as(z)
# Compute VQ loss
if self.legacy:
loss = F.mse_loss(z_q.detach(), z) + self.beta * F.mse_loss(z_q, z.detach())
else:
loss = self.beta * F.mse_loss(z_q.detach(), z) + F.mse_loss(z_q, z.detach())
# Optional entropy loss
entropy_aux = self.entropy_loss(d) if self.training else None
# Straight-through estimator
z_q = z + (z_q - z).detach()
# Reshape to original
z_q = rearrange(z_q, "b h w c -> b c h w")
z = rearrange(z, "b h w c -> b c h w")
# Remap if needed
if self.remap is not None:
min_indices = min_indices.view(z.shape[0], -1)
min_indices = self.remap_to_used(min_indices).view(-1, 1)
if self.sane_index_shape:
min_indices = min_indices.view(z_q.shape[0], z_q.shape[2], z_q.shape[3])
return {
"quantized": z_q,
"quantization_loss": loss,
"entropy_loss": entropy_aux,
"indices": min_indices,
}
def get_codebook_entry(self, indices, shape):
if self.remap is not None:
indices = indices.view(shape[0], -1)
indices = self.unmap_to_all(indices).view(-1)
z_q = self.embedding(indices)
if shape is not None:
z_q = z_q.view(shape).permute(0, 3, 1, 2).contiguous()
return z_q
|