(self, hidden_size, num_layers, ffn_kernel_size=9, dropout=None, num_heads=2,
use_pos_embed=True, use_last_norm=True, norm='ln', use_pos_embed_alpha=True)
| 250 | |
| 251 | class FFTBlocks(nn.Module): |
| 252 | def __init__(self, hidden_size, num_layers, ffn_kernel_size=9, dropout=None, num_heads=2, |
| 253 | use_pos_embed=True, use_last_norm=True, norm='ln', use_pos_embed_alpha=True): |
| 254 | super().__init__() |
| 255 | self.num_layers = num_layers |
| 256 | embed_dim = self.hidden_size = hidden_size |
| 257 | self.dropout = dropout if dropout is not None else hparams['dropout'] |
| 258 | self.use_pos_embed = use_pos_embed |
| 259 | self.use_last_norm = use_last_norm |
| 260 | if use_pos_embed: |
| 261 | self.max_source_positions = DEFAULT_MAX_TARGET_POSITIONS |
| 262 | self.padding_idx = 0 |
| 263 | self.pos_embed_alpha = nn.Parameter(torch.Tensor([1])) if use_pos_embed_alpha else 1 |
| 264 | self.embed_positions = SinusoidalPositionalEmbedding( |
| 265 | embed_dim, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS, |
| 266 | ) |
| 267 | |
| 268 | self.layers = nn.ModuleList([]) |
| 269 | self.layers.extend([ |
| 270 | TransformerEncoderLayer(self.hidden_size, self.dropout, |
| 271 | kernel_size=ffn_kernel_size, num_heads=num_heads) |
| 272 | for _ in range(self.num_layers) |
| 273 | ]) |
| 274 | if self.use_last_norm: |
| 275 | if norm == 'ln': |
| 276 | self.layer_norm = nn.LayerNorm(embed_dim) |
| 277 | elif norm == 'bn': |
| 278 | self.layer_norm = BatchNorm1dTBC(embed_dim) |
| 279 | else: |
| 280 | self.layer_norm = None |
| 281 | |
| 282 | def forward(self, x, padding_mask=None, attn_mask=None, return_hiddens=False): |
| 283 | """ |
nothing calls this directly
no test coverage detected