| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| """Positionwise feed forward layer definition.""" |
|
|
| import torch |
|
|
|
|
| class PositionwiseFeedForward(torch.nn.Module): |
| """Positionwise feed forward layer. |
| |
| FeedForward are appied on each position of the sequence. |
| The output dim is same with the input dim. |
| |
| Args: |
| idim (int): Input dimenstion. |
| hidden_units (int): The number of hidden units. |
| dropout_rate (float): Dropout rate. |
| activation (torch.nn.Module): Activation function |
| """ |
|
|
| def __init__(self, |
| idim: int, |
| hidden_units: int, |
| dropout_rate: float, |
| activation: torch.nn.Module = torch.nn.ReLU(), |
| adaptive_scale: bool = False, |
| init_weights: bool = False): |
| """Construct a PositionwiseFeedForward object.""" |
| super(PositionwiseFeedForward, self).__init__() |
| self.idim = idim |
| self.hidden_units = hidden_units |
| self.w_1 = torch.nn.Linear(idim, hidden_units) |
| self.activation = activation |
| self.dropout = torch.nn.Dropout(dropout_rate) |
| self.w_2 = torch.nn.Linear(hidden_units, idim) |
| self.ada_scale = None |
| self.ada_bias = None |
| self.adaptive_scale = adaptive_scale |
| self.ada_scale = torch.nn.Parameter(torch.ones([1, 1, idim]), |
| requires_grad=adaptive_scale) |
| self.ada_bias = torch.nn.Parameter(torch.zeros([1, 1, idim]), |
| requires_grad=adaptive_scale) |
| if init_weights: |
| self.init_weights() |
|
|
| def init_weights(self): |
| ffn1_max = self.idim**-0.5 |
| ffn2_max = self.hidden_units**-0.5 |
| torch.nn.init.uniform_(self.w_1.weight.data, -ffn1_max, ffn1_max) |
| torch.nn.init.uniform_(self.w_1.bias.data, -ffn1_max, ffn1_max) |
| torch.nn.init.uniform_(self.w_2.weight.data, -ffn2_max, ffn2_max) |
| torch.nn.init.uniform_(self.w_2.bias.data, -ffn2_max, ffn2_max) |
|
|
| def forward(self, xs: torch.Tensor) -> torch.Tensor: |
| """Forward function. |
| |
| Args: |
| xs: input tensor (B, L, D) |
| Returns: |
| output tensor, (B, L, D) |
| """ |
| if self.adaptive_scale: |
| xs = self.ada_scale * xs + self.ada_bias |
| return self.w_2(self.dropout(self.activation(self.w_1(xs)))) |
|
|
