| 106 | |
| 107 | class STARVectorQuantizer(nn.Module): |
| 108 | def __init__(self, num_embeddings, embedding_dim, commitment_cost): |
| 109 | super(STARVectorQuantizer, self).__init__() |
| 110 | |
| 111 | self._embedding_dim = embedding_dim |
| 112 | self._num_embeddings = num_embeddings |
| 113 | |
| 114 | self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim) |
| 115 | # self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings) |
| 116 | # pre_kmeans = torch.load("kmeans-centers-8192.pt") |
| 117 | # print("initializing the vq embedding with pre trained kmeans cluster", pre_kmeans.size()) |
| 118 | # self._embedding.weight.data.copy_(pre_kmeans) |
| 119 | self._commitment_cost = commitment_cost |
| 120 | self.data_initialized = 1 |
| 121 | |
| 122 | def forward(self, inputs): |
| 123 | # convert inputs from BCHW -> BHWC |