| 120 | self.data_initialized = 1 |
| 121 | |
| 122 | def forward(self, inputs): |
| 123 | # convert inputs from BCHW -> BHWC |
| 124 | # inputs = inputs.permute(0, 2, 3, 1).contiguous() |
| 125 | # print("vq inputs", inputs.size()) |
| 126 | input_shape = inputs.shape |
| 127 | # Flatten input |
| 128 | flat_input = inputs.view(-1, self._embedding_dim) |
| 129 | |
| 130 | # # L2 normalization before calculating the distances: |
| 131 | # normed_flat_input = flat_input / torch.linalg.norm(flat_input, dim=1, keepdim=True) |
| 132 | # normed_emb = self._embedding.weight / torch.linalg.norm(self._embedding.weight, dim=1, keepdim=True) |
| 133 | |
| 134 | # Following Andrej Karpathy's attempt to fix index collapse: |
| 135 | # https://github.com/karpathy/deep-vector-quantization/blob/main/dvq/model/quantize.py |
| 136 | if self.training and self.data_initialized ==0: |
| 137 | print("run kmeans") |
| 138 | rp = torch.randperm(flat_input.size(0)) |
| 139 | kd = kmeans2(flat_input[rp].data.cpu().numpy(), self._num_embeddings, minit="points") |
| 140 | self._embedding.weight.data.copy_(torch.from_numpy(kd[0])) |
| 141 | self.data_initialized = 1 |
| 142 | |
| 143 | # Calculate distances |
| 144 | # print(flat_input.size(), self._embedding.weight.size()) |
| 145 | distances = (torch.sum(flat_input**2, dim=1, keepdim=True) |
| 146 | + torch.sum(self._embedding.weight**2, dim=1, keepdim=True).t() |
| 147 | - 2 * torch.matmul(flat_input, self._embedding.weight.t())) |
| 148 | # distances = (torch.sum(normed_flat_input**2, dim=1, keepdim=True) |
| 149 | # + torch.sum(normed_emb**2, dim=1) |
| 150 | # - 2 * torch.matmul(normed_flat_input, normed_emb.t())) |
| 151 | |
| 152 | # Encoding |
| 153 | encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1) |
| 154 | encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device) |
| 155 | encodings.scatter_(1, encoding_indices, 1) |
| 156 | # print("encoding_indices", encoding_indices) |
| 157 | # print("encoding_indices", encoding_indices.size()) |
| 158 | # Quantize and unflatten |
| 159 | quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape) |
| 160 | # print("quantized shape", quantized.size()) |
| 161 | # Loss |
| 162 | e_latent_loss = F.mse_loss(quantized.detach(), inputs) |
| 163 | q_latent_loss = F.mse_loss(quantized, inputs.detach()) |
| 164 | loss = q_latent_loss + self._commitment_cost * e_latent_loss |
| 165 | |
| 166 | quantized = inputs + (quantized - inputs).detach() |
| 167 | avg_probs = torch.mean(encodings, dim=0) |
| 168 | # print(avg_probs) |
| 169 | # print("avg_probs", avg_probs.size()) |
| 170 | perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) |
| 171 | |
| 172 | # convert quantized from BHWC -> BCHW |
| 173 | return loss, quantized, perplexity, encodings |
| 174 | |
| 175 | |
| 176 | class STARVectorQuantizerEMA(nn.Module): |