add files

funasr_local/models/encoder/opennmt_encoders/ci_scorers.py

@@ -0,0 +1,38 @@
+import torch
+from torch.nn import functional as F
+
+
+class DotScorer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+            self,
+            xs_pad: torch.Tensor,
+            spk_emb: torch.Tensor,
+    ):
+        # xs_pad: B, T, D
+        # spk_emb: B, N, D
+        scores = torch.matmul(xs_pad, spk_emb.transpose(1, 2))
+        return scores
+
+    def convert_tf2torch(self, var_dict_tf, var_dict_torch):
+        return {}
+
+
+class CosScorer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+            self,
+            xs_pad: torch.Tensor,
+            spk_emb: torch.Tensor,
+    ):
+        # xs_pad: B, T, D
+        # spk_emb: B, N, D
+        scores = F.cosine_similarity(xs_pad.unsqueeze(2), spk_emb.unsqueeze(1), dim=-1)
+        return scores
+
+    def convert_tf2torch(self, var_dict_tf, var_dict_torch):
+        return {}

funasr_local/models/encoder/opennmt_encoders/conv_encoder.py

@@ -0,0 +1,277 @@
+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
+from typeguard import check_argument_types
+import numpy as np
+from funasr_local.modules.nets_utils import make_pad_mask
+from funasr_local.modules.layer_norm import LayerNorm
+from funasr_local.models.encoder.abs_encoder import AbsEncoder
+import math
+from funasr_local.modules.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",
+    ):
+        assert check_argument_types()
+        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
+
+    def gen_tf2torch_map_dict(self):
+        tensor_name_prefix_torch = self.tf2torch_tensor_name_prefix_torch
+        tensor_name_prefix_tf = self.tf2torch_tensor_name_prefix_tf
+        map_dict_local = {
+            # torch: conv1d.weight in "out_channel in_channel kernel_size"
+            # tf   : conv1d.weight in "kernel_size in_channel out_channel"
+            # torch: linear.weight in "out_channel in_channel"
+            # tf   :  dense.weight in "in_channel out_channel"
+            "{}.cnn_a.0.conv1d.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/cnn_a/conv1d/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+            "{}.cnn_a.0.conv1d.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/cnn_a/conv1d/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+
+            "{}.cnn_a.layeridx.conv1d.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/cnn_a/conv1d_layeridx/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+            "{}.cnn_a.layeridx.conv1d.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/cnn_a/conv1d_layeridx/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+        }
+        if self.out_units is not None:
+            # add output layer
+            map_dict_local.update({
+                "{}.conv_out.weight".format(tensor_name_prefix_torch):
+                    {"name": "{}/cnn_a/conv1d_{}/kernel".format(tensor_name_prefix_tf, self.num_layers),
+                     "squeeze": None,
+                     "transpose": (2, 1, 0),
+                     },  # tf: (1, 256, 256) -> torch: (256, 256, 1)
+                "{}.conv_out.bias".format(tensor_name_prefix_torch):
+                    {"name": "{}/cnn_a/conv1d_{}/bias".format(tensor_name_prefix_tf, self.num_layers),
+                     "squeeze": None,
+                     "transpose": None,
+                     },  # tf: (256,) -> torch: (256,)
+            })
+
+        return map_dict_local
+
+    def convert_tf2torch(self,
+                         var_dict_tf,
+                         var_dict_torch,
+                         ):
+
+        map_dict = self.gen_tf2torch_map_dict()
+
+        var_dict_torch_update = dict()
+        for name in sorted(var_dict_torch.keys(), reverse=False):
+            if name.startswith(self.tf2torch_tensor_name_prefix_torch):
+                # process special (first and last) layers
+                if name in map_dict:
+                    name_tf = map_dict[name]["name"]
+                    data_tf = var_dict_tf[name_tf]
+                    if map_dict[name]["squeeze"] is not None:
+                        data_tf = np.squeeze(data_tf, axis=map_dict[name]["squeeze"])
+                    if map_dict[name]["transpose"] is not None:
+                        data_tf = np.transpose(data_tf, map_dict[name]["transpose"])
+                    data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                    assert var_dict_torch[name].size() == data_tf.size(), \
+                        "{}, {}, {} != {}".format(name, name_tf,
+                                                  var_dict_torch[name].size(), data_tf.size())
+                    var_dict_torch_update[name] = data_tf
+                    logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                        name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                    ))
+                # process general layers
+                else:
+                    # self.tf2torch_tensor_name_prefix_torch may include ".", solve this case
+                    names = name.replace(self.tf2torch_tensor_name_prefix_torch, "todo").split('.')
+                    layeridx = int(names[2])
+                    name_q = name.replace(".{}.".format(layeridx), ".layeridx.")
+                    if name_q in map_dict.keys():
+                        name_v = map_dict[name_q]["name"]
+                        name_tf = name_v.replace("layeridx", "{}".format(layeridx))
+                        data_tf = var_dict_tf[name_tf]
+                        if map_dict[name_q]["squeeze"] is not None:
+                            data_tf = np.squeeze(data_tf, axis=map_dict[name_q]["squeeze"])
+                        if map_dict[name_q]["transpose"] is not None:
+                            data_tf = np.transpose(data_tf, map_dict[name_q]["transpose"])
+                        data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                        assert var_dict_torch[name].size() == data_tf.size(), \
+                            "{}, {}, {} != {}".format(name, name_tf,
+                                                      var_dict_torch[name].size(), data_tf.size())
+                        var_dict_torch_update[name] = data_tf
+                        logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                            name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                        ))
+                    else:
+                        logging.warning("{} is missed from tf checkpoint".format(name))
+
+        return var_dict_torch_update
+

funasr_local/models/encoder/opennmt_encoders/fsmn_encoder.py

@@ -0,0 +1,335 @@
+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
+from typeguard import check_argument_types
+import numpy as np
+from funasr_local.modules.nets_utils import make_pad_mask
+from funasr_local.modules.layer_norm import LayerNorm
+from funasr_local.models.encoder.abs_encoder import AbsEncoder
+import math
+from funasr_local.modules.repeat import repeat
+from funasr_local.modules.multi_layer_conv import FsmnFeedForward
+
+
+class FsmnBlock(torch.nn.Module):
+    def __init__(
+            self,
+            n_feat,
+            dropout_rate,
+            kernel_size,
+            fsmn_shift=0,
+    ):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout_rate)
+        self.fsmn_block = nn.Conv1d(n_feat, n_feat, kernel_size, stride=1,
+                                    padding=0, groups=n_feat, bias=False)
+        # padding
+        left_padding = (kernel_size - 1) // 2
+        if fsmn_shift > 0:
+            left_padding = left_padding + fsmn_shift
+        right_padding = kernel_size - 1 - left_padding
+        self.pad_fn = nn.ConstantPad1d((left_padding, right_padding), 0.0)
+
+    def forward(self, inputs, mask, mask_shfit_chunk=None):
+        b, t, d = inputs.size()
+        if mask is not None:
+            mask = torch.reshape(mask, (b, -1, 1))
+            if mask_shfit_chunk is not None:
+                mask = mask * mask_shfit_chunk
+
+        inputs = inputs * mask
+        x = inputs.transpose(1, 2)
+        x = self.pad_fn(x)
+        x = self.fsmn_block(x)
+        x = x.transpose(1, 2)
+        x = x + inputs
+        x = self.dropout(x)
+        return x * mask
+
+
+class EncoderLayer(torch.nn.Module):
+    def __init__(
+            self,
+            in_size,
+            size,
+            feed_forward,
+            fsmn_block,
+            dropout_rate=0.0
+    ):
+        super().__init__()
+        self.in_size = in_size
+        self.size = size
+        self.ffn = feed_forward
+        self.memory = fsmn_block
+        self.dropout = nn.Dropout(dropout_rate)
+
+    def forward(
+            self,
+            xs_pad: torch.Tensor,
+            mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # xs_pad in Batch, Time, Dim
+
+        context = self.ffn(xs_pad)[0]
+        memory = self.memory(context, mask)
+
+        memory = self.dropout(memory)
+        if self.in_size == self.size:
+            return memory + xs_pad, mask
+
+        return memory, mask
+
+
+class FsmnEncoder(AbsEncoder):
+    """Encoder using Fsmn
+      """
+
+    def __init__(self,
+                 in_units,
+                 filter_size,
+                 fsmn_num_layers,
+                 dnn_num_layers,
+                 num_memory_units=512,
+                 ffn_inner_dim=2048,
+                 dropout_rate=0.0,
+                 shift=0,
+                 position_encoder=None,
+                 sample_rate=1,
+                 out_units=None,
+                 tf2torch_tensor_name_prefix_torch="post_net",
+                 tf2torch_tensor_name_prefix_tf="EAND/post_net"
+                 ):
+        """Initializes the parameters of the encoder.
+
+        Args:
+          filter_size: the total order of memory block
+          fsmn_num_layers: The number of fsmn layers.
+          dnn_num_layers: The number of dnn layers
+          num_units: The number of memory units.
+          ffn_inner_dim: The number of units of the inner linear transformation
+            in the feed forward layer.
+          dropout_rate: The probability to drop units from the outputs.
+          shift: left padding, to control delay
+          position_encoder: The :class:`opennmt.layers.position.PositionEncoder` to
+            apply on inputs or ``None``.
+        """
+        super(FsmnEncoder, self).__init__()
+        self.in_units = in_units
+        self.filter_size = filter_size
+        self.fsmn_num_layers = fsmn_num_layers
+        self.dnn_num_layers = dnn_num_layers
+        self.num_memory_units = num_memory_units
+        self.ffn_inner_dim = ffn_inner_dim
+        self.dropout_rate = dropout_rate
+        self.shift = shift
+        if not isinstance(shift, list):
+            self.shift = [shift for _ in range(self.fsmn_num_layers)]
+        self.sample_rate = sample_rate
+        if not isinstance(sample_rate, list):
+            self.sample_rate = [sample_rate for _ in range(self.fsmn_num_layers)]
+        self.position_encoder = position_encoder
+        self.dropout = nn.Dropout(dropout_rate)
+        self.out_units = out_units
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+
+        self.fsmn_layers = repeat(
+            self.fsmn_num_layers,
+            lambda lnum: EncoderLayer(
+                in_units if lnum == 0 else num_memory_units,
+                num_memory_units,
+                FsmnFeedForward(
+                    in_units if lnum == 0 else num_memory_units,
+                    ffn_inner_dim,
+                    num_memory_units,
+                    1,
+                    dropout_rate
+                ),
+                FsmnBlock(
+                    num_memory_units,
+                    dropout_rate,
+                    filter_size,
+                    self.shift[lnum]
+                )
+            ),
+        )
+
+        self.dnn_layers = repeat(
+            dnn_num_layers,
+            lambda lnum: FsmnFeedForward(
+                num_memory_units,
+                ffn_inner_dim,
+                num_memory_units,
+                1,
+                dropout_rate,
+            )
+        )
+        if out_units is not None:
+            self.conv1d = nn.Conv1d(num_memory_units, out_units, 1, 1)
+
+    def output_size(self) -> int:
+        return self.num_memory_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)
+
+        inputs = self.dropout(inputs)
+        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+        inputs = self.fsmn_layers(inputs, masks)[0]
+        inputs = self.dnn_layers(inputs)[0]
+
+        if self.out_units is not None:
+            inputs = self.conv1d(inputs.transpose(1, 2)).transpose(1, 2)
+
+        return inputs, ilens, None
+
+    def gen_tf2torch_map_dict(self):
+        tensor_name_prefix_torch = self.tf2torch_tensor_name_prefix_torch
+        tensor_name_prefix_tf = self.tf2torch_tensor_name_prefix_tf
+        map_dict_local = {
+            # torch: conv1d.weight in "out_channel in_channel kernel_size"
+            # tf   : conv1d.weight in "kernel_size in_channel out_channel"
+            # torch: linear.weight in "out_channel in_channel"
+            # tf   :  dense.weight in "in_channel out_channel"
+            # for fsmn_layers
+            "{}.fsmn_layers.layeridx.ffn.norm.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/ffn/LayerNorm/beta".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.fsmn_layers.layeridx.ffn.norm.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/ffn/LayerNorm/gamma".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.fsmn_layers.layeridx.ffn.w_1.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/ffn/conv1d/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.fsmn_layers.layeridx.ffn.w_1.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/ffn/conv1d/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+            "{}.fsmn_layers.layeridx.ffn.w_2.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/ffn/conv1d_1/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+            "{}.fsmn_layers.layeridx.memory.fsmn_block.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/fsmn_layer_layeridx/memory/depth_conv_w".format(tensor_name_prefix_tf),
+                 "squeeze": 0,
+                 "transpose": (1, 2, 0),
+                 },  # (1, 31, 512, 1) -> (31, 512, 1) -> (512, 1, 31)
+
+            # for dnn_layers
+            "{}.dnn_layers.layeridx.norm.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/dnn_layer_layeridx/LayerNorm/beta".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.dnn_layers.layeridx.norm.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/dnn_layer_layeridx/LayerNorm/gamma".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.dnn_layers.layeridx.w_1.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/dnn_layer_layeridx/conv1d/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },
+            "{}.dnn_layers.layeridx.w_1.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/dnn_layer_layeridx/conv1d/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+            "{}.dnn_layers.layeridx.w_2.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/dnn_layer_layeridx/conv1d_1/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": (2, 1, 0),
+                 },
+
+        }
+        if self.out_units is not None:
+            # add output layer
+            map_dict_local.update({
+                "{}.conv1d.weight".format(tensor_name_prefix_torch):
+                    {"name": "{}/conv1d/kernel".format(tensor_name_prefix_tf),
+                     "squeeze": None,
+                     "transpose": (2, 1, 0),
+                     },
+                "{}.conv1d.bias".format(tensor_name_prefix_torch):
+                    {"name": "{}/conv1d/bias".format(tensor_name_prefix_tf),
+                     "squeeze": None,
+                     "transpose": None,
+                     },
+            })
+
+        return map_dict_local
+
+    def convert_tf2torch(self,
+                         var_dict_tf,
+                         var_dict_torch,
+                         ):
+
+        map_dict = self.gen_tf2torch_map_dict()
+
+        var_dict_torch_update = dict()
+        for name in sorted(var_dict_torch.keys(), reverse=False):
+            if name.startswith(self.tf2torch_tensor_name_prefix_torch):
+                # process special (first and last) layers
+                if name in map_dict:
+                    name_tf = map_dict[name]["name"]
+                    data_tf = var_dict_tf[name_tf]
+                    if map_dict[name]["squeeze"] is not None:
+                        data_tf = np.squeeze(data_tf, axis=map_dict[name]["squeeze"])
+                    if map_dict[name]["transpose"] is not None:
+                        data_tf = np.transpose(data_tf, map_dict[name]["transpose"])
+                    data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                    assert var_dict_torch[name].size() == data_tf.size(), \
+                        "{}, {}, {} != {}".format(name, name_tf,
+                                                  var_dict_torch[name].size(), data_tf.size())
+                    var_dict_torch_update[name] = data_tf
+                    logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                        name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                    ))
+                # process general layers
+                else:
+                    # self.tf2torch_tensor_name_prefix_torch may include ".", solve this case
+                    names = name.replace(self.tf2torch_tensor_name_prefix_torch, "todo").split('.')
+                    layeridx = int(names[2])
+                    name_q = name.replace(".{}.".format(layeridx), ".layeridx.")
+                    if name_q in map_dict.keys():
+                        name_v = map_dict[name_q]["name"]
+                        name_tf = name_v.replace("layeridx", "{}".format(layeridx))
+                        data_tf = var_dict_tf[name_tf]
+                        if map_dict[name_q]["squeeze"] is not None:
+                            data_tf = np.squeeze(data_tf, axis=map_dict[name_q]["squeeze"])
+                        if map_dict[name_q]["transpose"] is not None:
+                            data_tf = np.transpose(data_tf, map_dict[name_q]["transpose"])
+                        data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                        assert var_dict_torch[name].size() == data_tf.size(), \
+                            "{}, {}, {} != {}".format(name, name_tf,
+                                                      var_dict_torch[name].size(), data_tf.size())
+                        var_dict_torch_update[name] = data_tf
+                        logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                            name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                        ))
+                    else:
+                        logging.warning("{} is missed from tf checkpoint".format(name))
+
+        return var_dict_torch_update

funasr_local/models/encoder/opennmt_encoders/self_attention_encoder.py

@@ -0,0 +1,480 @@
+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 funasr_local.modules.streaming_utils.chunk_utilis import overlap_chunk
+from typeguard import check_argument_types
+import numpy as np
+from funasr_local.modules.nets_utils import make_pad_mask
+from funasr_local.modules.attention import MultiHeadSelfAttention, MultiHeadedAttentionSANM
+from funasr_local.modules.embedding import SinusoidalPositionEncoder
+from funasr_local.modules.layer_norm import LayerNorm
+from funasr_local.modules.multi_layer_conv import Conv1dLinear
+from funasr_local.modules.multi_layer_conv import MultiLayeredConv1d
+from funasr_local.modules.positionwise_feed_forward import (
+    PositionwiseFeedForward,  # noqa: H301
+)
+from funasr_local.modules.repeat import repeat
+from funasr_local.modules.subsampling import Conv2dSubsampling
+from funasr_local.modules.subsampling import Conv2dSubsampling2
+from funasr_local.modules.subsampling import Conv2dSubsampling6
+from funasr_local.modules.subsampling import Conv2dSubsampling8
+from funasr_local.modules.subsampling import TooShortUttError
+from funasr_local.modules.subsampling import check_short_utt
+from funasr_local.models.ctc import CTC
+from funasr_local.models.encoder.abs_encoder import AbsEncoder
+
+
+class EncoderLayer(nn.Module):
+    def __init__(
+        self,
+        in_size,
+        size,
+        self_attn,
+        feed_forward,
+        dropout_rate,
+        normalize_before=True,
+        concat_after=False,
+        stochastic_depth_rate=0.0,
+    ):
+        """Construct an EncoderLayer object."""
+        super(EncoderLayer, self).__init__()
+        self.self_attn = self_attn
+        self.feed_forward = feed_forward
+        self.norm1 = LayerNorm(in_size)
+        self.norm2 = LayerNorm(size)
+        self.dropout = nn.Dropout(dropout_rate)
+        self.in_size = in_size
+        self.size = size
+        self.normalize_before = normalize_before
+        self.concat_after = concat_after
+        if self.concat_after:
+            self.concat_linear = nn.Linear(size + size, size)
+        self.stochastic_depth_rate = stochastic_depth_rate
+        self.dropout_rate = dropout_rate
+
+    def forward(self, x, mask, cache=None, mask_att_chunk_encoder=None):
+        """Compute encoded features.
+
+        Args:
+            x_input (torch.Tensor): Input tensor (#batch, time, size).
+            mask (torch.Tensor): Mask tensor for the input (#batch, time).
+            cache (torch.Tensor): Cache tensor of the input (#batch, time - 1, size).
+
+        Returns:
+            torch.Tensor: Output tensor (#batch, time, size).
+            torch.Tensor: Mask tensor (#batch, time).
+
+        """
+        skip_layer = False
+        # with stochastic depth, residual connection `x + f(x)` becomes
+        # `x <- x + 1 / (1 - p) * f(x)` at training time.
+        stoch_layer_coeff = 1.0
+        if self.training and self.stochastic_depth_rate > 0:
+            skip_layer = torch.rand(1).item() < self.stochastic_depth_rate
+            stoch_layer_coeff = 1.0 / (1 - self.stochastic_depth_rate)
+
+        if skip_layer:
+            if cache is not None:
+                x = torch.cat([cache, x], dim=1)
+            return x, mask
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm1(x)
+
+        if self.concat_after:
+            x_concat = torch.cat((x, self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)), dim=-1)
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.concat_linear(x_concat)
+            else:
+                x = stoch_layer_coeff * self.concat_linear(x_concat)
+        else:
+            if self.in_size == self.size:
+                x = residual + stoch_layer_coeff * self.dropout(
+                    self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)
+                )
+            else:
+                x = stoch_layer_coeff * self.dropout(
+                    self.self_attn(x, mask, mask_att_chunk_encoder=mask_att_chunk_encoder)
+                )
+        if not self.normalize_before:
+            x = self.norm1(x)
+
+        residual = x
+        if self.normalize_before:
+            x = self.norm2(x)
+        x = residual + stoch_layer_coeff * self.dropout(self.feed_forward(x))
+        if not self.normalize_before:
+            x = self.norm2(x)
+
+        return x, mask, cache, mask_att_chunk_encoder
+
+
+class SelfAttentionEncoder(AbsEncoder):
+    """
+    Author: Speech Lab of DAMO Academy, Alibaba Group
+    Self attention encoder in OpenNMT framework
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int = 256,
+        attention_heads: int = 4,
+        linear_units: int = 2048,
+        num_blocks: int = 6,
+        dropout_rate: float = 0.1,
+        positional_dropout_rate: float = 0.1,
+        attention_dropout_rate: float = 0.0,
+        input_layer: Optional[str] = "conv2d",
+        pos_enc_class=SinusoidalPositionEncoder,
+        normalize_before: bool = True,
+        concat_after: bool = False,
+        positionwise_layer_type: str = "linear",
+        positionwise_conv_kernel_size: int = 1,
+        padding_idx: int = -1,
+        interctc_layer_idx: List[int] = [],
+        interctc_use_conditioning: bool = False,
+        tf2torch_tensor_name_prefix_torch: str = "encoder",
+        tf2torch_tensor_name_prefix_tf: str = "seq2seq/encoder",
+        out_units=None,
+    ):
+        assert check_argument_types()
+        super().__init__()
+        self._output_size = output_size
+
+        if input_layer == "linear":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Linear(input_size, output_size),
+                torch.nn.LayerNorm(output_size),
+                torch.nn.Dropout(dropout_rate),
+                torch.nn.ReLU(),
+                pos_enc_class(output_size, positional_dropout_rate),
+            )
+        elif input_layer == "conv2d":
+            self.embed = Conv2dSubsampling(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d2":
+            self.embed = Conv2dSubsampling2(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d6":
+            self.embed = Conv2dSubsampling6(input_size, output_size, dropout_rate)
+        elif input_layer == "conv2d8":
+            self.embed = Conv2dSubsampling8(input_size, output_size, dropout_rate)
+        elif input_layer == "embed":
+            self.embed = torch.nn.Sequential(
+                torch.nn.Embedding(input_size, output_size, padding_idx=padding_idx),
+                SinusoidalPositionEncoder(),
+            )
+        elif input_layer is None:
+            if input_size == output_size:
+                self.embed = None
+            else:
+                self.embed = torch.nn.Linear(input_size, output_size)
+        elif input_layer == "pe":
+            self.embed = SinusoidalPositionEncoder()
+        elif input_layer == "null":
+            self.embed = None
+        else:
+            raise ValueError("unknown input_layer: " + input_layer)
+        self.normalize_before = normalize_before
+        if positionwise_layer_type == "linear":
+            positionwise_layer = PositionwiseFeedForward
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d":
+            positionwise_layer = MultiLayeredConv1d
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        elif positionwise_layer_type == "conv1d-linear":
+            positionwise_layer = Conv1dLinear
+            positionwise_layer_args = (
+                output_size,
+                linear_units,
+                positionwise_conv_kernel_size,
+                dropout_rate,
+            )
+        else:
+            raise NotImplementedError("Support only linear or conv1d.")
+
+        self.encoders = repeat(
+            num_blocks,
+            lambda lnum: EncoderLayer(
+                output_size,
+                output_size,
+                MultiHeadSelfAttention(
+                    attention_heads,
+                    output_size,
+                    output_size,
+                    attention_dropout_rate,
+                ),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ) if lnum > 0 else EncoderLayer(
+                input_size,
+                output_size,
+                MultiHeadSelfAttention(
+                    attention_heads,
+                    input_size if input_layer == "pe" or input_layer == "null" else output_size,
+                    output_size,
+                    attention_dropout_rate,
+                ),
+                positionwise_layer(*positionwise_layer_args),
+                dropout_rate,
+                normalize_before,
+                concat_after,
+            ),
+        )
+        if self.normalize_before:
+            self.after_norm = LayerNorm(output_size)
+
+        self.interctc_layer_idx = interctc_layer_idx
+        if len(interctc_layer_idx) > 0:
+            assert 0 < min(interctc_layer_idx) and max(interctc_layer_idx) < num_blocks
+        self.interctc_use_conditioning = interctc_use_conditioning
+        self.conditioning_layer = None
+        self.dropout = nn.Dropout(dropout_rate)
+        self.tf2torch_tensor_name_prefix_torch = tf2torch_tensor_name_prefix_torch
+        self.tf2torch_tensor_name_prefix_tf = tf2torch_tensor_name_prefix_tf
+        self.out_units = out_units
+        if out_units is not None:
+            self.output_linear = nn.Linear(output_size, out_units)
+
+    def output_size(self) -> int:
+        return self._output_size
+
+    def forward(
+        self,
+        xs_pad: torch.Tensor,
+        ilens: torch.Tensor,
+        prev_states: torch.Tensor = None,
+        ctc: CTC = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        """Embed positions in tensor.
+
+        Args:
+            xs_pad: input tensor (B, L, D)
+            ilens: input length (B)
+            prev_states: Not to be used now.
+        Returns:
+            position embedded tensor and mask
+        """
+        masks = (~make_pad_mask(ilens)[:, None, :]).to(xs_pad.device)
+        xs_pad = xs_pad * self.output_size()**0.5
+        if self.embed is None:
+            xs_pad = xs_pad
+        elif (
+            isinstance(self.embed, Conv2dSubsampling)
+            or isinstance(self.embed, Conv2dSubsampling2)
+            or isinstance(self.embed, Conv2dSubsampling6)
+            or isinstance(self.embed, Conv2dSubsampling8)
+        ):
+            short_status, limit_size = check_short_utt(self.embed, xs_pad.size(1))
+            if short_status:
+                raise TooShortUttError(
+                    f"has {xs_pad.size(1)} frames and is too short for subsampling "
+                    + f"(it needs more than {limit_size} frames), return empty results",
+                    xs_pad.size(1),
+                    limit_size,
+                )
+            xs_pad, masks = self.embed(xs_pad, masks)
+        else:
+            xs_pad = self.embed(xs_pad)
+
+        xs_pad = self.dropout(xs_pad)
+        # encoder_outs = self.encoders0(xs_pad, masks)
+        # xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        intermediate_outs = []
+        if len(self.interctc_layer_idx) == 0:
+            encoder_outs = self.encoders(xs_pad, masks)
+            xs_pad, masks = encoder_outs[0], encoder_outs[1]
+        else:
+            for layer_idx, encoder_layer in enumerate(self.encoders):
+                encoder_outs = encoder_layer(xs_pad, masks)
+                xs_pad, masks = encoder_outs[0], encoder_outs[1]
+
+                if layer_idx + 1 in self.interctc_layer_idx:
+                    encoder_out = xs_pad
+
+                    # intermediate outputs are also normalized
+                    if self.normalize_before:
+                        encoder_out = self.after_norm(encoder_out)
+
+                    intermediate_outs.append((layer_idx + 1, encoder_out))
+
+                    if self.interctc_use_conditioning:
+                        ctc_out = ctc.softmax(encoder_out)
+                        xs_pad = xs_pad + self.conditioning_layer(ctc_out)
+
+        if self.normalize_before:
+            xs_pad = self.after_norm(xs_pad)
+
+        if self.out_units is not None:
+            xs_pad = self.output_linear(xs_pad)
+        olens = masks.squeeze(1).sum(1)
+        if len(intermediate_outs) > 0:
+            return (xs_pad, intermediate_outs), olens, None
+        return xs_pad, olens, None
+
+    def gen_tf2torch_map_dict(self):
+        tensor_name_prefix_torch = self.tf2torch_tensor_name_prefix_torch
+        tensor_name_prefix_tf = self.tf2torch_tensor_name_prefix_tf
+        map_dict_local = {
+            # cicd
+            # torch: conv1d.weight in "out_channel in_channel kernel_size"
+            # tf   : conv1d.weight in "kernel_size in_channel out_channel"
+            # torch: linear.weight in "out_channel in_channel"
+            # tf   :  dense.weight in "in_channel out_channel"
+            "{}.encoders.layeridx.norm1.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/LayerNorm/gamma".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            "{}.encoders.layeridx.norm1.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/LayerNorm/beta".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            "{}.encoders.layeridx.self_attn.linear_q_k_v.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/conv1d/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": 0,
+                 "transpose": (1, 0),
+                 },  # (768,256),(1,256,768)
+            "{}.encoders.layeridx.self_attn.linear_q_k_v.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/conv1d/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (768,),(768,)
+            "{}.encoders.layeridx.self_attn.linear_out.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/conv1d_1/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": 0,
+                 "transpose": (1, 0),
+                 },  # (256,256),(1,256,256)
+            "{}.encoders.layeridx.self_attn.linear_out.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/multi_head/conv1d_1/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            # ffn
+            "{}.encoders.layeridx.norm2.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/LayerNorm/gamma".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            "{}.encoders.layeridx.norm2.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/LayerNorm/beta".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            "{}.encoders.layeridx.feed_forward.w_1.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/conv1d/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": 0,
+                 "transpose": (1, 0),
+                 },  # (1024,256),(1,256,1024)
+            "{}.encoders.layeridx.feed_forward.w_1.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/conv1d/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (1024,),(1024,)
+            "{}.encoders.layeridx.feed_forward.w_2.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/conv1d_1/kernel".format(tensor_name_prefix_tf),
+                 "squeeze": 0,
+                 "transpose": (1, 0),
+                 },  # (256,1024),(1,1024,256)
+            "{}.encoders.layeridx.feed_forward.w_2.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/layer_layeridx/ffn/conv1d_1/bias".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            # out norm
+            "{}.after_norm.weight".format(tensor_name_prefix_torch):
+                {"name": "{}/LayerNorm/gamma".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+            "{}.after_norm.bias".format(tensor_name_prefix_torch):
+                {"name": "{}/LayerNorm/beta".format(tensor_name_prefix_tf),
+                 "squeeze": None,
+                 "transpose": None,
+                 },  # (256,),(256,)
+        }
+        if self.out_units is not None:
+            map_dict_local.update({
+                "{}.output_linear.weight".format(tensor_name_prefix_torch):
+                    {"name": "{}/conv1d/kernel".format(tensor_name_prefix_tf),
+                     "squeeze": 0,
+                     "transpose": (1, 0),
+                     },
+                "{}.output_linear.bias".format(tensor_name_prefix_torch):
+                    {"name": "{}/conv1d/bias".format(tensor_name_prefix_tf),
+                     "squeeze": None,
+                     "transpose": None,
+                     },  # (256,),(256,)
+            })
+    
+        return map_dict_local
+
+    def convert_tf2torch(self,
+                         var_dict_tf,
+                         var_dict_torch,
+                         ):
+        
+        map_dict = self.gen_tf2torch_map_dict()
+    
+        var_dict_torch_update = dict()
+        for name in sorted(var_dict_torch.keys(), reverse=False):
+            if name.startswith(self.tf2torch_tensor_name_prefix_torch):
+                # process special (first and last) layers
+                if name in map_dict:
+                    name_tf = map_dict[name]["name"]
+                    data_tf = var_dict_tf[name_tf]
+                    data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                    if map_dict[name]["squeeze"] is not None:
+                        data_tf = np.squeeze(data_tf, axis=map_dict[name]["squeeze"])
+                    if map_dict[name]["transpose"] is not None:
+                        data_tf = np.transpose(data_tf, map_dict[name]["transpose"])
+                    assert var_dict_torch[name].size() == data_tf.size(), \
+                        "{}, {}, {} != {}".format(name, name_tf,
+                                                  var_dict_torch[name].size(), data_tf.size())
+                    var_dict_torch_update[name] = data_tf
+                    logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                        name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                    ))
+                # process general layers
+                else:
+                    # self.tf2torch_tensor_name_prefix_torch may include ".", solve this case
+                    names = name.replace(self.tf2torch_tensor_name_prefix_torch, "todo").split('.')
+                    layeridx = int(names[2])
+                    name_q = name.replace(".{}.".format(layeridx), ".layeridx.")
+                    if name_q in map_dict.keys():
+                        name_v = map_dict[name_q]["name"]
+                        name_tf = name_v.replace("layeridx", "{}".format(layeridx))
+                        data_tf = var_dict_tf[name_tf]
+                        if map_dict[name_q]["squeeze"] is not None:
+                            data_tf = np.squeeze(data_tf, axis=map_dict[name_q]["squeeze"])
+                        if map_dict[name_q]["transpose"] is not None:
+                            data_tf = np.transpose(data_tf, map_dict[name_q]["transpose"])
+                        data_tf = torch.from_numpy(data_tf).type(torch.float32).to("cpu")
+                        assert var_dict_torch[name].size() == data_tf.size(), \
+                            "{}, {}, {} != {}".format(name, name_tf,
+                                                      var_dict_torch[name].size(), data_tf.size())
+                        var_dict_torch_update[name] = data_tf
+                        logging.info("torch tensor: {}, {}, loading from tf tensor: {}, {}".format(
+                            name, data_tf.size(), name_tf, var_dict_tf[name_tf].shape
+                        ))
+                    else:
+                        logging.warning("{} is missed from tf checkpoint".format(name))
+
+        return var_dict_torch_update