| 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 |