espnet2.asr.encoder.tcn_encoder.TCNEncoder

source

class espnet2.asr.encoder.tcn_encoder.TCNEncoder(input_size: int, output_size: int = 128, bottleneck_size: int = 64, hidden_size: int = 128, kernel_size: int = 3, num_blocks: int = 5, num_repeats: int = 2, norm_type: str = 'cLN', causal: bool | None = None, pre_output_nonlinear: str = 'prelu', , future_taps: int = 0, look_ahead: int = 0)

Bases: AbsLookAheadEncoder

Conv-TasNet TCN encoder whose look-ahead is chosen at run time.

• Parameters:
  • input_size – Input feature dimension.
  • output_size – Output dimension.
  • bottleneck_size – Channels between the blocks (B).
  • hidden_size – Channels inside each block (H).
  • kernel_size – Depthwise taps up to the current frame (P): the current frame and P - 1 past taps.
  • num_blocks – Blocks per repeat, at dilations 1, 2, …, 2**(X - 1) (X).
  • num_repeats – Number of repeats (R).
  • future_taps – Depthwise taps after the current frame that a layer can use. 0 makes a causal encoder. Keyword-only, like look_ahead.
  • look_ahead – Look-ahead in frames to start with; set_look_ahead() changes it.
  • norm_type – “cLN” (layer norm over channels, per frame) or “BN”. “gLN” and “GN” (over all frames) only with causal=False, for configurations of the previous encoder: every output then reads the whole sequence, so the encoder has no bounded look-ahead (the look-ahead cannot be changed, receptive_field and forward_streaming raise) and a padded utterance in a batch is normalized with the padding, as in the previous encoder.
  • pre_output_nonlinear – Nonlinearity before the output projection: “linear”, “prelu”, “relu”, “tanh” or “sigmoid”.
  • causal – The configuration of the previous TCNEncoder (a wrapper around espnet2.enh.layers.tcn.TemporalConvNet), where kernel_size counted all the taps. None: not used. True: future_taps 0 (the same network). False: the centred kernel of kernel_size taps (odd), i.e. (kernel_size + 1) // 2 taps up to the current frame and (kernel_size - 1) // 2 after it, at the full look-ahead. Give it without future_taps and look_ahead. Checkpoints of the previous encoder load: their state-dict keys are renamed.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(xs_pad: Tensor, ilens: Tensor, prev_states: Tensor | None = None) → Tuple[Tensor, Tensor, Tensor | None]

Encode a padded batch.

• Parameters:
  • xs_pad – Input features (B, T, input_size).
  • ilens – Input lengths (B,).
  • prev_states – Not used.
• Returns: Encoded features (B, T, output_size), output lengths (B,), None.

forward_streaming(xs: Tensor, state: List[Tensor] | None = None) → Tuple[Tensor, List[Tensor]]

Encode the next frames of a stream online (look-ahead 0).

Feeding a sequence in chunks of any sizes gives forward()’s output on the whole sequence at look-ahead 0, whatever look-ahead is set. With norm_type “BN” the encoder must be in eval mode. Every layer but the depthwise convolutions works per frame, so the state is the last (kernel_size - 1) * dilation depthwise inputs of each block: receptive_field - 1 frames in all, whatever the stream length.

• Parameters:
  • xs – New input frames (B, k, input_size), no padding.
  • state – The state returned by the previous call; None at the start of a stream.
• Returns: Encoded new frames (B, k, output_size) and the state for the next call.

future_taps_per_block() → List[int]

Number of future taps each block uses at the current look-ahead.

property look_ahead : int

Current look-ahead, in input frames.

property max_look_ahead : int

Largest look-ahead, in input frames.

output_size() → int

property past_context : int

Number of past input frames each output frame reads.

property receptive_field : int

Number of input frames an output frame reads at look-ahead 0.

The current frame and past_context frames before it. A streaming caller can rerun the encoder on each chunk’s new frames together with the receptive_field - 1 frames before them and get the whole-recording output for the new frames.

• Raises:RuntimeError – With gLN/GN, where every output reads every frame.

set_look_ahead(frames: int)

Make each output frame t read input frames up to t + frames.

The look-ahead goes to the blocks with the largest dilation first: a block takes as many future taps as fit, up to future_taps, and the rest goes to the next one.

• Raises:ValueError – frames is negative, larger than max_look_ahead, or cannot be made from the blocks’ future taps.

upgrade_state_dict(state_dict: dict, prefix: str = '')

Convert a state dict of the previous TCNEncoder in place.

That encoder held espnet2.enh.layers.tcn.TemporalConvNet as tcn: network.0 input norm, network.1 bottleneck, network.2.<r>.<x> blocks (net.0 1x1 conv, net.1 PReLU, net.2 norm, net.3.net depthwise conv, [Chomp1d when causal,] PReLU, norm, 1x1 conv), then the output, after a nonlinearity module unless it was linear. Its keys are renamed, and a depthwise kernel with fewer taps than this encoder’s gets zero future taps appended: a causal checkpoint then initializes an encoder with future taps and gives the same output at every look-ahead until the future taps are trained. load_state_dict() calls it; call it directly before code that compares keys or shapes first (e.g. espnet2.torch_utils.load_pretrained_model). Other state dicts are left as they are.