| 174 | |
| 175 | |
| 176 | class STARVectorQuantizerEMA(nn.Module): |
| 177 | def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5): |
| 178 | super(STARVectorQuantizerEMA, self).__init__() |
| 179 | |
| 180 | self._embedding_dim = embedding_dim |
| 181 | self._num_embeddings = num_embeddings |
| 182 | |
| 183 | self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim) |
| 184 | self._embedding.weight.data.normal_() |
| 185 | self._commitment_cost = commitment_cost |
| 186 | |
| 187 | self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings)) |
| 188 | self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim)) |
| 189 | # pre_kmeans = torch.load("kmeans-centers-8192.pt") |
| 190 | # print("initializing the vq embedding with pre trained kmeans cluster", pre_kmeans.size()) |
| 191 | # self._ema_w.data.copy_(pre_kmeans) |
| 192 | # self._ema_w.data.normal_() |
| 193 | self.data_initialized = 0 |
| 194 | self._decay = decay |
| 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) |