Spaces:
Running
on
A10G
Running
on
A10G
| from typing import Optional | |
| import torch | |
| import torch.nn.functional as F | |
| from torch import nn | |
| from .wavenet import WaveNet | |
| class ReferenceEncoder(WaveNet): | |
| def __init__( | |
| self, | |
| input_channels: Optional[int] = None, | |
| output_channels: Optional[int] = None, | |
| residual_channels: int = 512, | |
| residual_layers: int = 20, | |
| dilation_cycle: Optional[int] = 4, | |
| num_heads: int = 8, | |
| latent_len: int = 4, | |
| ): | |
| super().__init__( | |
| input_channels=input_channels, | |
| residual_channels=residual_channels, | |
| residual_layers=residual_layers, | |
| dilation_cycle=dilation_cycle, | |
| ) | |
| self.head_dim = residual_channels // num_heads | |
| self.num_heads = num_heads | |
| self.latent_len = latent_len | |
| self.latent = nn.Parameter(torch.zeros(1, self.latent_len, residual_channels)) | |
| self.q = nn.Linear(residual_channels, residual_channels, bias=True) | |
| self.kv = nn.Linear(residual_channels, residual_channels * 2, bias=True) | |
| self.q_norm = nn.LayerNorm(self.head_dim) | |
| self.k_norm = nn.LayerNorm(self.head_dim) | |
| self.proj = nn.Linear(residual_channels, residual_channels) | |
| self.proj_drop = nn.Dropout(0.1) | |
| self.norm = nn.LayerNorm(residual_channels) | |
| self.mlp = nn.Sequential( | |
| nn.Linear(residual_channels, residual_channels * 4), | |
| nn.SiLU(), | |
| nn.Linear(residual_channels * 4, residual_channels), | |
| ) | |
| self.output_projection_attn = nn.Linear(residual_channels, output_channels) | |
| torch.nn.init.trunc_normal_(self.latent, std=0.02) | |
| self.apply(self.init_weights) | |
| def init_weights(self, m): | |
| if isinstance(m, nn.Linear): | |
| torch.nn.init.trunc_normal_(m.weight, std=0.02) | |
| if m.bias is not None: | |
| torch.nn.init.constant_(m.bias, 0) | |
| def forward(self, x, attn_mask=None): | |
| x = super().forward(x).mT | |
| B, N, C = x.shape | |
| # Calculate mask | |
| if attn_mask is not None: | |
| assert attn_mask.shape == (B, N) and attn_mask.dtype == torch.bool | |
| attn_mask = attn_mask[:, None, None, :].expand( | |
| B, self.num_heads, self.latent_len, N | |
| ) | |
| q_latent = self.latent.expand(B, -1, -1) | |
| q = ( | |
| self.q(q_latent) | |
| .reshape(B, self.latent_len, self.num_heads, self.head_dim) | |
| .transpose(1, 2) | |
| ) | |
| kv = ( | |
| self.kv(x) | |
| .reshape(B, N, 2, self.num_heads, self.head_dim) | |
| .permute(2, 0, 3, 1, 4) | |
| ) | |
| k, v = kv.unbind(0) | |
| q, k = self.q_norm(q), self.k_norm(k) | |
| x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) | |
| x = x.transpose(1, 2).reshape(B, self.latent_len, C) | |
| x = self.proj(x) | |
| x = self.proj_drop(x) | |
| x = x + self.mlp(self.norm(x)) | |
| x = self.output_projection_attn(x) | |
| x = x.mean(1) | |
| return x | |
| if __name__ == "__main__": | |
| with torch.autocast(device_type="cpu", dtype=torch.bfloat16): | |
| model = ReferenceEncoder( | |
| input_channels=128, | |
| output_channels=64, | |
| residual_channels=384, | |
| residual_layers=20, | |
| dilation_cycle=4, | |
| num_heads=8, | |
| ) | |
| x = torch.randn(4, 128, 64) | |
| mask = torch.ones(4, 64, dtype=torch.bool) | |
| y = model(x, mask) | |
| print(y.shape) | |
| loss = F.mse_loss(y, torch.randn(4, 64)) | |
| loss.backward() | |