| from typing import List |
| from typing import Optional |
| from typing import Sequence |
| from typing import Tuple |
| from typing import Union |
| import logging |
| import torch |
| import torch.nn as nn |
| from torch.nn import functional as F |
| import numpy as np |
| from funasr.models.transformer.utils.nets_utils import make_pad_mask |
| from funasr.models.transformer.layer_norm import LayerNorm |
| from funasr.models.encoder.abs_encoder import AbsEncoder |
| import math |
| from funasr.models.transformer.utils.repeat import repeat |
|
|
|
|
| class EncoderLayer(nn.Module): |
| def __init__( |
| self, |
| input_units, |
| num_units, |
| kernel_size=3, |
| activation="tanh", |
| stride=1, |
| include_batch_norm=False, |
| residual=False, |
| ): |
| super().__init__() |
| left_padding = math.ceil((kernel_size - stride) / 2) |
| right_padding = kernel_size - stride - left_padding |
| self.conv_padding = nn.ConstantPad1d((left_padding, right_padding), 0.0) |
| self.conv1d = nn.Conv1d( |
| input_units, |
| num_units, |
| kernel_size, |
| stride, |
| ) |
| self.activation = self.get_activation(activation) |
| if include_batch_norm: |
| self.bn = nn.BatchNorm1d(num_units, momentum=0.99, eps=1e-3) |
| self.residual = residual |
| self.include_batch_norm = include_batch_norm |
| self.input_units = input_units |
| self.num_units = num_units |
| self.stride = stride |
|
|
| @staticmethod |
| def get_activation(activation): |
| if activation == "tanh": |
| return nn.Tanh() |
| else: |
| return nn.ReLU() |
|
|
| def forward(self, xs_pad, ilens=None): |
| outputs = self.conv1d(self.conv_padding(xs_pad)) |
| if self.residual and self.stride == 1 and self.input_units == self.num_units: |
| outputs = outputs + xs_pad |
|
|
| if self.include_batch_norm: |
| outputs = self.bn(outputs) |
|
|
| |
| return self.activation(outputs), ilens |
|
|
|
|
| class ConvEncoder(AbsEncoder): |
| """ |
| Author: Speech Lab of DAMO Academy, Alibaba Group |
| Convolution encoder in OpenNMT framework |
| """ |
|
|
| def __init__( |
| self, |
| num_layers, |
| input_units, |
| num_units, |
| kernel_size=3, |
| dropout_rate=0.3, |
| position_encoder=None, |
| activation="tanh", |
| auxiliary_states=True, |
| out_units=None, |
| out_norm=False, |
| out_residual=False, |
| include_batchnorm=False, |
| regularization_weight=0.0, |
| stride=1, |
| tf2torch_tensor_name_prefix_torch: str = "speaker_encoder", |
| tf2torch_tensor_name_prefix_tf: str = "EAND/speaker_encoder", |
| ): |
| super().__init__() |
| self._output_size = num_units |
|
|
| self.num_layers = num_layers |
| self.input_units = input_units |
| self.num_units = num_units |
| self.kernel_size = kernel_size |
| self.dropout_rate = dropout_rate |
| self.position_encoder = position_encoder |
| self.out_units = out_units |
| self.auxiliary_states = auxiliary_states |
| self.out_norm = out_norm |
| self.activation = activation |
| self.out_residual = out_residual |
| self.include_batch_norm = include_batchnorm |
| self.regularization_weight = regularization_weight |
| self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch |
| self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf |
| if isinstance(stride, int): |
| self.stride = [stride] * self.num_layers |
| else: |
| self.stride = stride |
| self.downsample_rate = 1 |
| for s in self.stride: |
| self.downsample_rate *= s |
|
|
| self.dropout = nn.Dropout(dropout_rate) |
| self.cnn_a = repeat( |
| self.num_layers, |
| lambda lnum: EncoderLayer( |
| input_units if lnum == 0 else num_units, |
| num_units, |
| kernel_size, |
| activation, |
| self.stride[lnum], |
| include_batchnorm, |
| residual=True if lnum > 0 else False, |
| ), |
| ) |
|
|
| if self.out_units is not None: |
| left_padding = math.ceil((kernel_size - stride) / 2) |
| right_padding = kernel_size - stride - left_padding |
| self.out_padding = nn.ConstantPad1d((left_padding, right_padding), 0.0) |
| self.conv_out = nn.Conv1d( |
| num_units, |
| out_units, |
| kernel_size, |
| ) |
|
|
| if self.out_norm: |
| self.after_norm = LayerNorm(out_units) |
|
|
| def output_size(self) -> int: |
| return self.num_units |
|
|
| def forward( |
| self, |
| xs_pad: torch.Tensor, |
| ilens: torch.Tensor, |
| prev_states: torch.Tensor = None, |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: |
|
|
| inputs = xs_pad |
| if self.position_encoder is not None: |
| inputs = self.position_encoder(inputs) |
|
|
| if self.dropout_rate > 0: |
| inputs = self.dropout(inputs) |
|
|
| outputs, _ = self.cnn_a(inputs.transpose(1, 2), ilens) |
|
|
| if self.out_units is not None: |
| outputs = self.conv_out(self.out_padding(outputs)) |
|
|
| outputs = outputs.transpose(1, 2) |
| if self.out_norm: |
| outputs = self.after_norm(outputs) |
|
|
| if self.out_residual: |
| outputs = outputs + inputs |
|
|
| return outputs, ilens, None |
|
|
