(
self,
input_dim: int,
embed_dim: int,
num_layers: int,
num_heads: int,
max_len: int,
mask_num_chunks: int = 6,
mask_chunk_width: int = 50,
slice_section_length: int = 20,
slice_overlap: int = 10,
dropout: float = 0.1,
norm_first: bool = False,
)
| 12 | |
| 13 | class SpecFormer(L.LightningModule): |
| 14 | def __init__( |
| 15 | self, |
| 16 | input_dim: int, |
| 17 | embed_dim: int, |
| 18 | num_layers: int, |
| 19 | num_heads: int, |
| 20 | max_len: int, |
| 21 | mask_num_chunks: int = 6, |
| 22 | mask_chunk_width: int = 50, |
| 23 | slice_section_length: int = 20, |
| 24 | slice_overlap: int = 10, |
| 25 | dropout: float = 0.1, |
| 26 | norm_first: bool = False, |
| 27 | ): |
| 28 | super().__init__() |
| 29 | self.save_hyperparameters() |
| 30 | |
| 31 | self.data_embed = nn.Linear(input_dim, embed_dim) |
| 32 | self.position_embed = nn.Embedding(max_len, embed_dim) |
| 33 | self.dropout = nn.Dropout(dropout) |
| 34 | self.blocks = nn.ModuleList( |
| 35 | [ |
| 36 | TransformerBlock( |
| 37 | embedding_dim=embed_dim, |
| 38 | num_heads=num_heads, |
| 39 | causal=False, |
| 40 | dropout=dropout, |
| 41 | bias=True, |
| 42 | ) |
| 43 | for _ in range(num_layers) |
| 44 | ] |
| 45 | ) |
| 46 | self.final_layernorm = LayerNorm(embed_dim, bias=True) |
| 47 | self.head = nn.Linear(embed_dim, input_dim, bias=True) |
| 48 | |
| 49 | self._reset_parameters_datapt() |
| 50 | |
| 51 | def forward(self, x: Tensor) -> torch.Tensor: |
| 52 | """Forward pass through the model.""" |
nothing calls this directly
no test coverage detected