| 17 | |
| 18 | |
| 19 | class SslAdaptor(nn.Module): |
| 20 | def __init__( |
| 21 | self, |
| 22 | in_dim: int, |
| 23 | embed_dim: int, |
| 24 | out_dim: int, |
| 25 | num_layers: int, |
| 26 | num_heads: int, |
| 27 | ffn_dim: int = None, |
| 28 | attn_dropout: float = 0.0, |
| 29 | dropout: float = 0.0, |
| 30 | ): |
| 31 | super().__init__() |
| 32 | self.in_dim = in_dim |
| 33 | self.embed_dim = embed_dim |
| 34 | self.dropout = dropout |
| 35 | # Input Projection |
| 36 | self.in_proj = nn.Linear(in_dim, embed_dim) |
| 37 | # Transformer |
| 38 | self.layers = nn.ModuleList( |
| 39 | [ |
| 40 | WhisperEncoderLayer( |
| 41 | embed_dim, num_heads, ffn_dim, attn_dropout, dropout |
| 42 | ) |
| 43 | for _ in range(num_layers) |
| 44 | ] |
| 45 | ) |
| 46 | # Output norm |
| 47 | self.layer_norm = nn.LayerNorm(embed_dim) |
| 48 | # Output projection |
| 49 | self.out_proj = nn.Linear(embed_dim, out_dim) |
| 50 | # Init weight |
| 51 | self.apply(self._init_weights) |
| 52 | |
| 53 | def forward( |
| 54 | self, |
| 55 | hidden_states: torch.Tensor, |
| 56 | hidden_length: torch.Tensor, |
| 57 | ): |
| 58 | # Downsampling |
| 59 | hidden_states = self.in_proj(hidden_states) |
| 60 | # Transformer |
| 61 | attention_mask = make_nonpad_mask(hidden_length).unsqueeze(1) # (b, 1, t) |
| 62 | for layer in self.layers: |
| 63 | hidden_states = layer(hidden_states, attention_mask) |
| 64 | hidden_states = self.layer_norm(hidden_states) |
| 65 | hidden_states = self.out_proj(hidden_states) |
| 66 | return hidden_states, hidden_length |
| 67 | |
| 68 | def _init_weights(self, module): |
| 69 | std = 0.02 |
| 70 | if isinstance(module, (nn.Linear, nn.Conv1d)): |
| 71 | module.weight.data.normal_(mean=0.0, std=std) |
| 72 | if module.bias is not None: |
| 73 | module.bias.data.zero_() |
| 74 | elif isinstance(module, nn.Embedding): |
| 75 | module.weight.data.normal_(mean=0.0, std=std) |
| 76 | if module.padding_idx is not None: |