75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
# Related papers:
|
|
# Shengkui Zhao, Trung Hieu Nguyen, Bin Ma, “Monaural Speech Enhancement with Complex Convolutional
|
|
# Block Attention Module and Joint Time Frequency Losses”, ICASSP 2021.
|
|
# Shiliang Zhang, Ming Lei, Zhijie Yan, Lirong Dai, “Deep-FSMN for Large Vocabulary Continuous Speech
|
|
# Recognition “, arXiv:1803.05030, 2018.
|
|
|
|
from torch import nn
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models import MODELS, TorchModel
|
|
from modelscope.models.audio.ans.layers.activations import (RectifiedLinear,
|
|
Sigmoid)
|
|
from modelscope.models.audio.ans.layers.affine_transform import AffineTransform
|
|
from modelscope.models.audio.ans.layers.uni_deep_fsmn import UniDeepFsmn
|
|
from modelscope.utils.constant import Tasks
|
|
|
|
|
|
@MODELS.register_module(
|
|
Tasks.acoustic_noise_suppression, module_name=Models.speech_dfsmn_ans)
|
|
class DfsmnAns(TorchModel):
|
|
"""Denoise model with DFSMN.
|
|
|
|
Args:
|
|
model_dir (str): the model path.
|
|
fsmn_depth (int): the depth of deepfsmn
|
|
lorder (int):
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_dir: str,
|
|
fsmn_depth=9,
|
|
lorder=20,
|
|
*args,
|
|
**kwargs):
|
|
super().__init__(model_dir, *args, **kwargs)
|
|
self.lorder = lorder
|
|
self.linear1 = AffineTransform(120, 256)
|
|
self.relu = RectifiedLinear(256, 256)
|
|
repeats = [
|
|
UniDeepFsmn(256, 256, lorder, 256) for i in range(fsmn_depth)
|
|
]
|
|
self.deepfsmn = nn.Sequential(*repeats)
|
|
self.linear2 = AffineTransform(256, 961)
|
|
self.sig = Sigmoid(961, 961)
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Args:
|
|
input: fbank feature [batch_size,number_of_frame,feature_dimension]
|
|
|
|
Returns:
|
|
mask value [batch_size, number_of_frame, FFT_size/2+1]
|
|
"""
|
|
x1 = self.linear1(input)
|
|
x2 = self.relu(x1)
|
|
x3 = self.deepfsmn(x2)
|
|
x4 = self.linear2(x3)
|
|
x5 = self.sig(x4)
|
|
return x5
|
|
|
|
def to_kaldi_nnet(self):
|
|
re_str = ''
|
|
re_str += '<Nnet>\n'
|
|
re_str += self.linear1.to_kaldi_nnet()
|
|
re_str += self.relu.to_kaldi_nnet()
|
|
for dfsmn in self.deepfsmn:
|
|
re_str += dfsmn.to_kaldi_nnet()
|
|
re_str += self.linear2.to_kaldi_nnet()
|
|
re_str += self.sig.to_kaldi_nnet()
|
|
re_str += '</Nnet>\n'
|
|
|
|
return re_str
|