| 684 | |
| 685 | |
| 686 | class Quantizer_module(torch.nn.Module): |
| 687 | def __init__(self, n_e, e_dim): |
| 688 | super(Quantizer_module, self).__init__() |
| 689 | self.embedding = nn.Embedding(n_e, e_dim) |
| 690 | self.embedding.weight.data.uniform_(-1.0 / n_e, 1.0 / n_e) |
| 691 | |
| 692 | def forward(self, x): |
| 693 | d = ( |
| 694 | torch.sum(x**2, 1, keepdim=True) |
| 695 | + torch.sum(self.embedding.weight**2, 1) |
| 696 | - 2 * torch.matmul(x, self.embedding.weight.T) |
| 697 | ) |
| 698 | min_indicies = torch.argmin(d, 1) |
| 699 | z_q = self.embedding(min_indicies) |
| 700 | return z_q, min_indicies |
| 701 | |
| 702 | |
| 703 | class Quantizer(torch.nn.Module): |