    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: str = "conv2d",
        pos_enc_layer_type: str = "rel_pos",
        normalize_before: bool = True,
        static_chunk_size: int = 0,
        use_dynamic_chunk: bool = False,
        global_cmvn: torch.nn.Module = None,
        use_dynamic_left_chunk: bool = False,
        macaron_style: bool = True,
        activation_type: str = "swish",
        use_cnn_module: bool = True,
        cnn_module_kernel: int = 15,
        causal: bool = False,
        cnn_module_norm: str = "batch_norm",
        stride_layer_idx: Optional[Union[int, List[int]]] = 3,
        stride: Optional[Union[int, List[int]]] = 2,
        group_layer_idx: Optional[Union[int, List[int], tuple]] = (0, 1, 2, 3),
        group_size: int = 3,
        stride_kernel: bool = True,
        **kwargs,
    ):
        """Construct Efficient Conformer Encoder

        Args:
            input_size to use_dynamic_chunk, see in BaseEncoder
            macaron_style (bool): Whether to use macaron style for
                positionwise layer.
            activation_type (str): Encoder activation function type.
            use_cnn_module (bool): Whether to use convolution module.
            cnn_module_kernel (int): Kernel size of convolution module.
            causal (bool): whether to use causal convolution or not.
            stride_layer_idx (list): layer id with StrideConv, start from 0
            stride (list): stride size of each StrideConv in efficient conformer
            group_layer_idx (list): layer id with GroupedAttention, start from 0
            group_size (int): group size of every GroupedAttention layer
            stride_kernel (bool): default True. True: recompute cnn kernels with stride.
        """
        super().__init__()
        self._output_size = output_size

        if pos_enc_layer_type == "abs_pos":
            pos_enc_class = PositionalEncoding
        elif pos_enc_layer_type == "rel_pos":
            pos_enc_class = RelPositionalEncoding
        elif pos_enc_layer_type == "no_pos":
            pos_enc_class = NoPositionalEncoding
        else:
            raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type)

        if input_layer == "linear":
            subsampling_class = LinearNoSubsampling
        elif input_layer == "conv2d2":
            subsampling_class = Conv2dSubsampling2
        elif input_layer == "conv2d":
            subsampling_class = Conv2dSubsampling4
        elif input_layer == "conv2d6":
            subsampling_class = Conv2dSubsampling6
        elif input_layer == "conv2d8":
            subsampling_class = Conv2dSubsampling8
        else:
            raise ValueError("unknown input_layer: " + input_layer)

        logging.info(
            f"input_layer = {input_layer}, " f"subsampling_class = {subsampling_class}"
        )

        self.global_cmvn = global_cmvn
        self.embed = subsampling_class(
            input_size,
            output_size,
            dropout_rate,
            pos_enc_class(output_size, positional_dropout_rate),
        )
        self.input_layer = input_layer
        self.normalize_before = normalize_before
        self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5)
        self.static_chunk_size = static_chunk_size
        self.use_dynamic_chunk = use_dynamic_chunk
        self.use_dynamic_left_chunk = use_dynamic_left_chunk

        activation = get_activation(activation_type)
        self.num_blocks = num_blocks
        self.attention_heads = attention_heads
        self.cnn_module_kernel = cnn_module_kernel
        self.global_chunk_size = 0
        self.chunk_feature_map = 0

        # efficient conformer configs
        self.stride_layer_idx = (
            [stride_layer_idx] if type(stride_layer_idx) == int else stride_layer_idx
        )
        self.stride = [stride] if type(stride) == int else stride
        self.group_layer_idx = (
            [group_layer_idx] if type(group_layer_idx) == int else group_layer_idx
        )
        self.grouped_size = group_size  # group size of every GroupedAttention layer

        assert len(self.stride) == len(self.stride_layer_idx)
        self.cnn_module_kernels = [cnn_module_kernel]  # kernel size of each StridedConv
        for i in self.stride:
            if stride_kernel:
                self.cnn_module_kernels.append(self.cnn_module_kernels[-1] // i)
            else:
                self.cnn_module_kernels.append(self.cnn_module_kernels[-1])

        logging.info(
            f"stride_layer_idx= {self.stride_layer_idx}, "
            f"stride = {self.stride}, "
            f"cnn_module_kernel = {self.cnn_module_kernels}, "
            f"group_layer_idx = {self.group_layer_idx}, "
            f"grouped_size = {self.grouped_size}"
        )

        # feed-forward module definition
        positionwise_layer = PositionwiseFeedForward
        positionwise_layer_args = (
            output_size,
            linear_units,
            dropout_rate,
            activation,
        )
        # convolution module definition
        convolution_layer = ConvolutionModule

        # encoder definition
        index = 0
        layers = []
        for i in range(num_blocks):
            # self-attention module definition
            if i in self.group_layer_idx:
                encoder_selfattn_layer = GroupedRelPositionMultiHeadedAttention
                encoder_selfattn_layer_args = (
                    attention_heads,
                    output_size,
                    attention_dropout_rate,
                    self.grouped_size,
                )
            else:
                if pos_enc_layer_type == "no_pos":
                    encoder_selfattn_layer = MultiHeadedAttention
                else:
                    encoder_selfattn_layer = RelPositionMultiHeadedAttention
                encoder_selfattn_layer_args = (
                    attention_heads,
                    output_size,
                    attention_dropout_rate,
                )

            # conformer module definition
            if i in self.stride_layer_idx:
                # conformer block with downsampling
                convolution_layer_args_stride = (
                    output_size,
                    self.cnn_module_kernels[index],
                    activation,
                    cnn_module_norm,
                    causal,
                    True,
                    self.stride[index],
                )
                layers.append(
                    StrideConformerEncoderLayer(
                        output_size,
                        encoder_selfattn_layer(*encoder_selfattn_layer_args),
                        positionwise_layer(*positionwise_layer_args),
                        (
                            positionwise_layer(*positionwise_layer_args)
                            if macaron_style
                            else None
                        ),
                        (
                            convolution_layer(*convolution_layer_args_stride)
                            if use_cnn_module
                            else None
                        ),
                        torch.nn.AvgPool1d(
                            kernel_size=self.stride[index],
                            stride=self.stride[index],
                            padding=0,
                            ceil_mode=True,
                            count_include_pad=False,
                        ),  # pointwise_conv_layer
                        dropout_rate,
                        normalize_before,
                    )
                )
                index = index + 1
            else:
                # conformer block
                convolution_layer_args_normal = (
                    output_size,
                    self.cnn_module_kernels[index],
                    activation,
                    cnn_module_norm,
                    causal,
                )
                layers.append(
                    ConformerEncoderLayer(
                        output_size,
                        encoder_selfattn_layer(*encoder_selfattn_layer_args),
                        positionwise_layer(*positionwise_layer_args),
                        (
                            positionwise_layer(*positionwise_layer_args)
                            if macaron_style
                            else None
                        ),
                        (
                            convolution_layer(*convolution_layer_args_normal)
                            if use_cnn_module
                            else None
                        ),
                        dropout_rate,
                        normalize_before,
                    )
                )

        self.encoders = torch.nn.ModuleList(layers)
Construct Efficient Conformer Encoder

Args:
input_size to use_dynamic_chunk, see in BaseEncoder
macaron_style (bool): Whether to use macaron style for
positionwise layer.
activation_type (str): Encoder activation function type.
use_cnn_module (bool): Whether to use convolution module.
cnn_module_kernel (int): Kernel size of convolution module.
causal (bool): whether to use causal convolution or not.
stride_layer_idx (list): layer id with StrideConv, start from 0
stride (list): stride size of each StrideConv in efficient conformer
group_layer_idx (list): layer id with GroupedAttention, start from 0
group_size (int): group size of every GroupedAttention layer
stride_kernel (bool): default True. True: recompute cnn kernels with stride.
modules/wenet_extractor/efficient_conformer/encoder.py