| 195 | self._epsilon = epsilon |
| 196 | |
| 197 | def forward(self, inputs): |
| 198 | # convert inputs from BCHW -> BHWC |
| 199 | # inputs = inputs.permute(0, 2, 3, 1).contiguous() |
| 200 | input_shape = inputs.shape |
| 201 | |
| 202 | # Flatten input |
| 203 | flat_input = inputs.view(-1, self._embedding_dim) |
| 204 | |
| 205 | # Following Andrej Karpathy's attempt to fix index collapse: |
| 206 | # https://github.com/karpathy/deep-vector-quantization/blob/main/dvq/model/quantize.py |
| 207 | if self.training and self.data_initialized ==0: |
| 208 | print("run kmeans") |
| 209 | rp = torch.randperm(flat_input.size(0)) |
| 210 | kd = kmeans2(flat_input[rp].data.cpu().numpy(), self._num_embeddings, minit="points") |
| 211 | self._ema_w.data.copy_(torch.from_numpy(kd[0])) |
| 212 | self.data_initialized = 1 |
| 213 | |
| 214 | # Calculate distances |
| 215 | distances = (torch.sum(flat_input**2, dim=1, keepdim=True) |
| 216 | + torch.sum(self._embedding.weight**2, dim=1) |
| 217 | - 2 * torch.matmul(flat_input, self._embedding.weight.t())) |
| 218 | |
| 219 | # Encoding |
| 220 | encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1) |
| 221 | encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device) |
| 222 | encodings.scatter_(1, encoding_indices, 1) |
| 223 | |
| 224 | # Quantize and unflatten |
| 225 | quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape) |
| 226 | |
| 227 | # Use EMA to update the embedding vectors |
| 228 | if self.training: |
| 229 | self._ema_cluster_size = self._ema_cluster_size * self._decay + \ |
| 230 | (1 - self._decay) * torch.sum(encodings, 0) |
| 231 | |
| 232 | # Laplace smoothing of the cluster size |
| 233 | n = torch.sum(self._ema_cluster_size.data) |
| 234 | self._ema_cluster_size = ( |
| 235 | (self._ema_cluster_size + self._epsilon) |
| 236 | / (n + self._num_embeddings * self._epsilon) * n) |
| 237 | |
| 238 | dw = torch.matmul(encodings.t(), flat_input) |
| 239 | self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw) |
| 240 | |
| 241 | self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1)) |
| 242 | |
| 243 | # Loss |
| 244 | e_latent_loss = F.mse_loss(quantized.detach(), inputs) |
| 245 | loss = self._commitment_cost * e_latent_loss |
| 246 | |
| 247 | # Straight Through Estimator |
| 248 | quantized = inputs + (quantized - inputs).detach() |
| 249 | avg_probs = torch.mean(encodings, dim=0) |
| 250 | perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) |
| 251 | |
| 252 | # convert quantized from BHWC -> BCHW |
| 253 | return loss, quantized, perplexity, encodings |
| 254 | |