Spaces:
Running
Running
# This module is from [WeNet](https://github.com/wenet-e2e/wenet). | |
# ## Citations | |
# ```bibtex | |
# @inproceedings{yao2021wenet, | |
# title={WeNet: Production oriented Streaming and Non-streaming End-to-End Speech Recognition Toolkit}, | |
# author={Yao, Zhuoyuan and Wu, Di and Wang, Xiong and Zhang, Binbin and Yu, Fan and Yang, Chao and Peng, Zhendong and Chen, Xiaoyu and Xie, Lei and Lei, Xin}, | |
# booktitle={Proc. Interspeech}, | |
# year={2021}, | |
# address={Brno, Czech Republic }, | |
# organization={IEEE} | |
# } | |
# @article{zhang2022wenet, | |
# title={WeNet 2.0: More Productive End-to-End Speech Recognition Toolkit}, | |
# author={Zhang, Binbin and Wu, Di and Peng, Zhendong and Song, Xingchen and Yao, Zhuoyuan and Lv, Hang and Xie, Lei and Yang, Chao and Pan, Fuping and Niu, Jianwei}, | |
# journal={arXiv preprint arXiv:2203.15455}, | |
# year={2022} | |
# } | |
# | |
"""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(), | |
): | |
"""Construct a PositionwiseFeedForward object.""" | |
super(PositionwiseFeedForward, self).__init__() | |
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) | |
def forward(self, xs: torch.Tensor) -> torch.Tensor: | |
"""Forward function. | |
Args: | |
xs: input tensor (B, L, D) | |
Returns: | |
output tensor, (B, L, D) | |
""" | |
return self.w_2(self.dropout(self.activation(self.w_1(xs)))) | |