(
self,
in_dim: int,
embed_dim: int,
out_dim: int,
num_layers: int,
num_heads: int,
ffn_dim: int = None,
attn_dropout: float = 0.0,
dropout: float = 0.0,
)
| 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, |
nothing calls this directly
no test coverage detected