tcn¶

Classes

`Conv2dStem`(n_mels[, channels, n_layers, ...])	Local time-frequency processing in front of the 1D trunk.
`PositionalEncoding`(channels)	Additive sinusoidal positional encoding (Vaswani et al., 2017).
`SelfAttentionBlock`(channels, n_heads)	One transformer-encoder-style block: self-attention sublayer, then a feedforward sublayer, each wrapped in its own residual connection and LayerNorm.
`TCNTempoNet`([n_mels, sample_rate, ...])	Dilated TCN for tempo regression (Davies & Böck, 2019), or per-frame beat-phase detection.

class Conv2dStem(n_mels, channels=16, n_layers=3, freq_pool=3)[source]¶

Bases: Module

Local time-frequency processing in front of the 1D trunk.

TCNTempoNet’s first operation used to be Conv1d(n_mels, channels, kernel_size=1): at every frame, a fixed linear mixture of all mel bands, after which the band axis is gone and every remaining layer convolves over time only. Nothing in the model ever saw a time-frequency neighbourhood.

That is the wrong first operation for beat tracking, for two reasons (plans/08_rethinking_the_approach.md §2.1 has the long version):

A 1x1 mixer is frequency-absolute; onsets are frequency-relative. It learns one weight per band, applied identically at every frame — so it can learn “these bands matter”, but not “energy rose in whichever band it rose in”. A kick drum and a walking bass note are the same event shape at different absolute frequencies, and a mixer must spend separate output channels on each register to detect the same thing twice. A Conv2d shares one kernel across frequency and gets that equivariance for free, which is the reason to put audio on a log-frequency axis at all.
Mixing before differencing lets onsets cancel. Onset strength is a difference across time within a band. With the bands summed first, a band rising and another falling by the same weighted amount produces a flat mixture, and the event is gone before any layer could see it. That is exactly a harmonic change with no percussive attack — the dominant downbeat cue wherever no drum marks the bar.

Every published tracker does local spectro-temporal processing first (madmom’s TCN opens with 3x3 convolutions and frequency max-pooling; Beat This! uses frequency-wise partial attention), and both reach their numbers with a weak decoder or none at all, which is what points at the front end.

Shape, with the defaults and n_mels=128:

(B, 128, T)                 log-mel, as the trunk used to receive it
(B, 1, 128, T)              band axis promoted to a spatial axis
(B, 16, 128, T)             Conv2d 3x3 + BN + GELU
(B, 16,  42, T)             MaxPool2d((3, 1)) — frequency only
(B, 16,  42, T)             Conv2d 3x3 + BN + GELU
(B, 16,  14, T)             MaxPool2d((3, 1))
(B, 16,  14, T)             Conv2d 3x3 + BN + GELU
(B, 224, T)                 frequency folded into channels

Time is never pooled. The frame rate is the output resolution, and the trunk’s dilations are what buy context — pooling time here would spend precision the 70 ms evaluation tolerance cannot afford. The 3x3 kernels do widen the receptive field by 2 * n_layers frames, which is negligible beside the trunk’s 1 + 2 * sum(dilations).

Parameters:

n_mels (int) – Number of input mel bands.
channels (int) – Feature maps per 2D layer. 16 is madmom-scale; the stem is meant to be cheap next to the trunk, not to hold capacity.
n_layers (int) – Number of Conv2d blocks. Frequency is pooled after every block except the last, so n_layers=3 pools twice.
freq_pool (int) – Frequency pooling factor per pool.

Raises:

ValueError – If n_layers is below 1, or if the pooling schedule would leave fewer than one frequency bin.

forward(x)[source]¶

Parameters:: x (Tensor) – Normalised log-mel, shape (B, n_mels, T).
Returns:: (B, out_channels, T) — T unchanged.
Return type:: Tensor

n_freq¶: Frequency bins surviving the pooling schedule.

out_channels¶: Channel count the trunk’s input_proj must expect.

class PositionalEncoding(channels)[source]¶

Bases: Module

Additive sinusoidal positional encoding (Vaswani et al., 2017).

Self-attention has no built-in notion of frame order (unlike convolution or recurrence), so this injects one: each position gets a fixed sin/cos pattern that varies by position and by channel pair, added directly to the input.

Parameters:: channels (int) – Channel width. Must be even — sin/cos are paired per two channels.

forward(x)[source]¶

Parameters:: x (Tensor) – (B, T, C)
Returns:: x plus the positional encoding, same shape.
Return type:: Tensor

class SelfAttentionBlock(channels, n_heads)[source]¶

Bases: Module

One transformer-encoder-style block: self-attention sublayer, then a feedforward sublayer, each wrapped in its own residual connection and LayerNorm. Lets every frame’s representation draw on every other frame in the sequence, unlike the TCN trunk’s fixed dilated-conv receptive field (see docs/beat_phase_context_ideas.md).

Parameters:

channels (int) – Channel width (attention embedding dim).
n_heads (int) – Number of attention heads.

forward(x)[source]¶

x -> multihead attention + res + layernorm -> h h -> MLP -> residual -> layernorm -> output

Parameters:: x (Tensor) – (B, T, C)
Return type:: Tensor

class TCNTempoNet(n_mels=128, sample_rate=22050, hop_length=512, channels=32, n_layers=8, dropout=0.3, n_outputs=1, frame_level=False, use_self_attention=False, n_attn_layers=1, n_attn_heads=4, conv2d_stem=False, stem_channels=16, stem_layers=3, stem_freq_pool=3, fixed_norm=False)[source]¶

Bases: Module

Dilated TCN for tempo regression (Davies & Böck, 2019), or per-frame beat-phase detection.

Applies a log-mel transform, projects to the TCN channel width, then runs a stack of dilated 1D residual convolutions with exponentially growing dilation (1, 2, 4, …, 2^(n_layers-1)).

Two output modes, controlled by frame_level:

frame_level=False (default): globally pools over time, then a small FC head produces scalar/bin regression or classification logits. Input: (B, 1, T) → Output: (B,) or (B, n_outputs).
frame_level=True: skips the pool; a 1x1 conv head produces per-frame logits instead (e.g. beat/one/last for beat-phase detection). Input: (B, 1, T) → Output: (B, n_outputs, T’) or (B, T’) if n_outputs == 1, where T’ is the mel transform’s frame count. Sigmoid is not applied — pair with BCEWithLogitsLoss downstream, matching the classification mode’s convention of returning raw logits.

Receptive field is 1 + (kernel_size - 1) * sum(dilations) frames — with the default schedule 1, 2, ..., 2^(n_layers-1) that is 2^(n_layers+1) - 1, so 511 frames ≈ 11.9 s at n_layers=8, hop_length=512. (This docstring used to quote kernel_size × (2^n_layers − 1), a loose upper bound ~1.5x the truth; see docs/beat_phase_context_ideas.md and plans/08 §1.2/§3.1.) The same trunk is shared between both modes, so this is unaffected by frame_level. A conv2d_stem adds 2 × stem_layers frames to that, which is noise beside it.

The receptive field is only real if the input is at least that long. Every trunk conv uses padding=dilation, so a layer whose dilation exceeds the input length has both off-centre taps in zero padding at every frame and collapses to a 1x1 conv. On a 16 s clip at hop_length=512 (689 frames) that is any layer past the ninth. See plans/08 §3.1.

conv2d_stem is the one structural option here. Off, the first operation is a 1x1 mix over mel bands and no layer ever sees a time-frequency neighbourhood; on, a small Conv2dStem runs first. See that class for why it exists.

Parameters:

n_mels (int) – Number of mel filterbanks.
sample_rate (int) – Audio sample rate used to build the mel transform.
hop_length (int) – Hop length for the mel transform. Controls temporal resolution (smaller = more frames per second). Defaults to 512 (≈43 fps at 22050 Hz).
channels (int) – Channel width for the TCN.
n_layers (int) – Number of dilated layers; dilation doubles per layer. Keep the dilation of the deepest layer (2^(n_layers-1) frames) below the input sequence length, or that layer only ever convolves padding.
dropout (float) – Dropout probability applied right before each head’s final Conv1d/Linear — the pooled regression head’s last Linear, the frame head’s 1x1 conv, or (when use_self_attention=True) the beat_head/phase_head 1x1 convs.
n_outputs (int) – Output dimension. In pooled mode, 1 for scalar regression, > 1 for classification over tempo bins. In frame-level mode, the number of per-frame target channels (e.g. 3 for beat/one/last).
frame_level (bool) – If True, produce per-frame outputs instead of pooling over time.
use_self_attention (bool) – Frame-level mode only. If True, splits the frame head in two: beat_head reads straight off the TCN trunk (unchanged, already accurate), while phase_head routes the remaining n_outputs - 1 channels through a positional encoding + a stack of SelfAttentionBlock, giving them context beyond the trunk’s fixed dilated-conv receptive field. Output channel order is always beat first, then the phase channels — (beat, one, last) for musicality.losses.beat_phase.beat_phase_loss(), or (beat, pos_1, ..., pos_G) for musicality.losses.beat_position.beat_position_loss(). See docs/beat_phase_context_ideas.md.
n_attn_layers (int) – Number of stacked SelfAttentionBlock in phase_head. Only used when use_self_attention=True.
n_attn_heads (int) – Attention heads per SelfAttentionBlock. Only used when use_self_attention=True.
conv2d_stem (bool) – Run a Conv2dStem between the mel and the trunk instead of projecting the raw bands. Defaults to False here so that checkpoints predating the stem reconstruct into identical parameter shapes from their own saved hyperparameters; the shipped frame-level backbones turn it on. Measured cost at their channels=32: 32,805 parameters to 40,773 — the stem itself is 4.9 k, the rest is input_proj widening from 128 to 224 inputs.
stem_channels (int) – Feature maps per stem layer. conv2d_stem only.
stem_layers (int) – Conv2d blocks in the stem; frequency is pooled after all but the last. conv2d_stem only.
stem_freq_pool (int) – Frequency pooling factor per stem pool. conv2d_stem only.
fixed_norm (bool) – Normalise the log-mel with frozen per-band statistics (norm_mean/norm_std, filled by fit_input_stats() at the start of training and saved in the checkpoint) instead of statistics taken over the input tensor. Off reproduces what every checkpoint up to v6 was trained with; on removes a train/inference mismatch, since the per-tensor statistics depend on whether the input is a 16 s crop or a whole track. plans/08 §2.1 item 4 has the measurement.

forward(wav)[source]¶

Parameters:: wav (Tensor)
Return type:: Tensor