(self, xin)
| 714 | self.embed_dim = embed_dim |
| 715 | |
| 716 | def forward(self, xin): |
| 717 | # B, C, T |
| 718 | B, C, T = xin.shape |
| 719 | xin = xin.transpose(1, 2) |
| 720 | x = xin.reshape(-1, self.embed_dim) |
| 721 | x = torch.split(x, self.embed_dim // self.n_code_groups, dim=-1) |
| 722 | min_indicies = [] |
| 723 | z_q = [] |
| 724 | for _x, m in zip(x, self.quantizer_modules): |
| 725 | _z_q, _min_indicies = m(_x) |
| 726 | z_q.append(_z_q) |
| 727 | min_indicies.append(_min_indicies) # B * T, |
| 728 | z_q = torch.cat(z_q, -1).reshape(xin.shape) |
| 729 | loss = 0.25 * torch.mean((z_q.detach() - xin) ** 2) + torch.mean( |
| 730 | (z_q - xin.detach()) ** 2 |
| 731 | ) |
| 732 | z_q = xin + (z_q - xin).detach() |
| 733 | z_q = z_q.transpose(1, 2) |
| 734 | codes = torch.stack(min_indicies, -1).reshape(B, T, self.n_code_groups) |
| 735 | return z_q, loss, codes.transpose(1, 2) |
| 736 | |
| 737 | def embed(self, x): |
| 738 | # idx: N, 4, T |
nothing calls this directly
no outgoing calls
no test coverage detected