funasr pages

游雁

2023-11-16 4ace5a95b052d338947fc88809a440ccd55cf6b4

 funasr/export/models/encoder/sanm_encoder.py

@@ -8,6 +8,7 @@
from funasr.export.models.modules.encoder_layer import EncoderLayerSANM as EncoderLayerSANM_export
from funasr.modules.positionwise_feed_forward import PositionwiseFeedForward
from funasr.export.models.modules.feedforward import PositionwiseFeedForward as PositionwiseFeedForward_export
from funasr.modules.embedding import StreamSinusoidalPositionEncoder


class SANMEncoder(nn.Module):
@@ -21,6 +22,8 @@
    ):
        super().__init__()
        self.embed = model.embed
        if isinstance(self.embed, StreamSinusoidalPositionEncoder):
            self.embed = None
        self.model = model
        self.feats_dim = feats_dim
        self._output_size = model._output_size
@@ -63,8 +66,10 @@
    def forward(self,
                speech: torch.Tensor,
                speech_lengths: torch.Tensor,
                online: bool = False 
                ):
        speech = speech * self._output_size ** 0.5
        if not online:
            speech = speech * self._output_size ** 0.5
        mask = self.make_pad_mask(speech_lengths)
        mask = self.prepare_mask(mask)
        if self.embed is None:
@@ -158,13 +163,14 @@
    def forward(self,
                speech: torch.Tensor,
                speech_lengths: torch.Tensor,
                vad_mask: torch.Tensor,
                vad_masks: torch.Tensor,
                sub_masks: torch.Tensor,
                ):
        speech = speech * self._output_size ** 0.5
        mask = self.make_pad_mask(speech_lengths)
        vad_masks = self.prepare_mask(mask, vad_masks)
        mask = self.prepare_mask(mask, sub_masks)
        vad_mask = self.prepare_mask(mask, vad_mask)
        
        if self.embed is None:
            xs_pad = speech
        else:
@@ -176,7 +182,7 @@
        # encoder_outs = self.model.encoders(xs_pad, mask)
        for layer_idx, encoder_layer in enumerate(self.model.encoders):
            if layer_idx == len(self.model.encoders) - 1:
                mask = vad_mask
                mask = vad_masks
            encoder_outs = encoder_layer(xs_pad, mask)
            xs_pad, masks = encoder_outs[0], encoder_outs[1]
        
@@ -187,26 +193,26 @@
    def get_output_size(self):
        return self.model.encoders[0].size
    
    def get_dummy_inputs(self):
        feats = torch.randn(1, 100, self.feats_dim)
        return (feats)
    
    def get_input_names(self):
        return ['feats']
    
    def get_output_names(self):
        return ['encoder_out', 'encoder_out_lens', 'predictor_weight']
    
    def get_dynamic_axes(self):
        return {
            'feats': {
                1: 'feats_length'
            },
            'encoder_out': {
                1: 'enc_out_length'
            },
            'predictor_weight': {
                1: 'pre_out_length'
            }
            
        }
    # def get_dummy_inputs(self):
    #     feats = torch.randn(1, 100, self.feats_dim)
    #     return (feats)
    #
    # def get_input_names(self):
    #     return ['feats']
    #
    # def get_output_names(self):
    #     return ['encoder_out', 'encoder_out_lens', 'predictor_weight']
    #
    # def get_dynamic_axes(self):
    #     return {
    #         'feats': {
    #             1: 'feats_length'
    #         },
    #         'encoder_out': {
    #             1: 'enc_out_length'
    #         },
    #         'predictor_weight': {
    #             1: 'pre_out_length'
    #         }
    #
    #     }