| 747 | |
| 748 | |
| 749 | class CodePredictor(nn.Module): |
| 750 | def __init__( |
| 751 | self, |
| 752 | hidden_channels, |
| 753 | filter_channels, |
| 754 | n_heads, |
| 755 | n_layers, |
| 756 | kernel_size, |
| 757 | p_dropout, |
| 758 | n_q=8, |
| 759 | dims=1024, |
| 760 | ssl_dim=768, |
| 761 | ): |
| 762 | super().__init__() |
| 763 | self.hidden_channels = hidden_channels |
| 764 | self.filter_channels = filter_channels |
| 765 | self.n_heads = n_heads |
| 766 | self.n_layers = n_layers |
| 767 | self.kernel_size = kernel_size |
| 768 | self.p_dropout = p_dropout |
| 769 | |
| 770 | self.vq_proj = nn.Conv1d(ssl_dim, hidden_channels, 1) |
| 771 | self.ref_enc = modules.MelStyleEncoder( |
| 772 | ssl_dim, style_vector_dim=hidden_channels |
| 773 | ) |
| 774 | |
| 775 | self.encoder = attentions.Encoder( |
| 776 | hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout |
| 777 | ) |
| 778 | |
| 779 | self.out_proj = nn.Conv1d(hidden_channels, (n_q - 1) * dims, 1) |
| 780 | self.n_q = n_q |
| 781 | self.dims = dims |
| 782 | |
| 783 | def forward(self, x, x_mask, refer, codes, infer=False): |
| 784 | x = x.detach() |
| 785 | x = self.vq_proj(x * x_mask) * x_mask |
| 786 | g = self.ref_enc(refer, x_mask) |
| 787 | x = x + g |
| 788 | x = self.encoder(x * x_mask, x_mask) |
| 789 | x = self.out_proj(x * x_mask) * x_mask |
| 790 | logits = x.reshape(x.shape[0], self.n_q - 1, self.dims, x.shape[-1]).transpose( |
| 791 | 2, 3 |
| 792 | ) |
| 793 | target = codes[1:].transpose(0, 1) |
| 794 | if not infer: |
| 795 | logits = logits.reshape(-1, self.dims) |
| 796 | target = target.reshape(-1) |
| 797 | loss = torch.nn.functional.cross_entropy(logits, target) |
| 798 | return loss |
| 799 | else: |
| 800 | _, top10_preds = torch.topk(logits, 10, dim=-1) |
| 801 | correct_top10 = torch.any(top10_preds == target.unsqueeze(-1), dim=-1) |
| 802 | top3_acc = 100 * torch.mean(correct_top10.float()).detach().cpu().item() |
| 803 | |
| 804 | print("Top-10 Accuracy:", top3_acc, "%") |
| 805 | |
| 806 | pred_codes = torch.argmax(logits, dim=-1) |
nothing calls this directly
no outgoing calls
no test coverage detected