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)
|
|
# add parenthesis for repeat module
|
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
|
