| 343 | |
| 344 | |
| 345 | class EncoderLayerSANM(nn.Module): |
| 346 | def __init__( |
| 347 | self, |
| 348 | in_size, |
| 349 | size, |
| 350 | self_attn, |
| 351 | feed_forward, |
| 352 | dropout_rate, |
| 353 | normalize_before=True, |
| 354 | concat_after=False, |
| 355 | stochastic_depth_rate=0.0, |
| 356 | ): |
| 357 | """Construct an EncoderLayer object.""" |
| 358 | super(EncoderLayerSANM, self).__init__() |
| 359 | self.self_attn = self_attn |
| 360 | self.feed_forward = feed_forward |
| 361 | self.norm1 = LayerNorm(in_size) |
| 362 | self.norm2 = LayerNorm(size) |
| 363 | self.dropout = nn.Dropout(dropout_rate) |
| 364 | self.in_size = in_size |
| 365 | self.size = size |
| 366 | self.normalize_before = normalize_before |
| 367 | self.concat_after = concat_after |
| 368 | if self.concat_after: |
| 369 | self.concat_linear = nn.Linear(size + size, size) |
| 370 | self.stochastic_depth_rate = stochastic_depth_rate |
| 371 | self.dropout_rate = dropout_rate |
| 372 | |
| 373 | def forward(self, x, mask, cache=None, mask_shfit_chunk=None, mask_att_chunk_encoder=None): |
| 374 | """Compute encoded features. |
| 375 | |
| 376 | Args: |
| 377 | x_input (torch.Tensor): Input tensor (#batch, time, size). |
| 378 | mask (torch.Tensor): Mask tensor for the input (#batch, time). |
| 379 | cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size). |
| 380 | |
| 381 | Returns: |
| 382 | torch.Tensor: Output tensor (#batch, time, size). |
| 383 | torch.Tensor: Mask tensor (#batch, time). |
| 384 | |
| 385 | """ |
| 386 | skip_layer = False |
| 387 | # with stochastic depth, residual connection `x + f(x)` becomes |
| 388 | # `x <- x + 1 / (1 - p) * f(x)` at training time. |
| 389 | stoch_layer_coeff = 1.0 |
| 390 | if self.training and self.stochastic_depth_rate > 0: |
| 391 | skip_layer = torch.rand(1).item() < self.stochastic_depth_rate |
| 392 | stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate) |
| 393 | |
| 394 | if skip_layer: |
| 395 | if cache is not None: |
| 396 | x = torch.cat([cache, x], dim=1) |
| 397 | return x, mask |
| 398 | |
| 399 | residual = x |
| 400 | if self.normalize_before: |
| 401 | x = self.norm1(x) |
| 402 |
no outgoing calls
no test coverage detected
searching dependent graphs…