"""BTC: the Bi-directional Transformer for Chord Recognition (ISMIR 2019). Inference-only port of Jonggwon Park's reference implementation (https://github.com/jayg996/BTC-ISMIR19, MIT licensed, see LICENSE-BTC), kept faithful enough to load the published checkpoints unchanged. Everything the checkpoint has a weight for is reproduced exactly, quirks included: * the layer norm divides by an *unbiased* standard deviation (`torch.std`), not the biased one `nn.LayerNorm` uses, so it can't be swapped for the built-in module; * the position-wise feed-forward applies its ReLU after *every* convolution, the last one included, so its output is non-negative; * attention is masked in both directions — one block sees only the past, its twin only the future — which is what makes the model "bi-directional". Dropout is left out (this only ever runs under `eval()`), and so is the unused LSTM the reference output layer carries: see `load_state_dict` below. """ import math import numpy as np import torch import torch.nn.functional as F from torch import nn def timing_signal(length: int, channels: int) -> torch.Tensor: """Sinusoidal position encoding, [1, length, channels]. The transformer's original formulation, as ported into BTC from tensor2tensor: `channels // 2` sine components followed by the same number of cosines, over timescales spanning 1 … 1e4. """ position = np.arange(length) num_timescales = channels // 2 log_increment = math.log(1.0e4) / (num_timescales - 1) inv_timescales = np.exp( np.arange(num_timescales, dtype=np.float64) * -log_increment ) scaled = np.expand_dims(position, 1) * np.expand_dims(inv_timescales, 0) signal = np.concatenate([np.sin(scaled), np.cos(scaled)], axis=1) signal = np.pad(signal, [[0, 0], [0, channels % 2]], "constant") return torch.from_numpy(signal.reshape([1, length, channels])).float() def causal_mask(length: int) -> torch.Tensor: """Additive attention mask, [1, 1, length, length], hiding future frames. Transposing the last two dimensions turns it into the mask that hides the past instead — which is exactly how the backward block is built. """ mask = np.triu(np.full([length, length], -np.inf), 1) return torch.from_numpy(mask).float().unsqueeze(0).unsqueeze(1) class LayerNorm(nn.Module): """Layer norm over the last dimension, with an unbiased standard deviation. `nn.LayerNorm` normalizes by the biased (population) standard deviation. BTC was trained against this variant, so the two are not interchangeable: swapping it in shifts every activation slightly and costs accuracy. """ def __init__(self, features: int, eps: float = 1e-6): super().__init__() self.gamma = nn.Parameter(torch.ones(features)) self.beta = nn.Parameter(torch.zeros(features)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: mean = x.mean(-1, keepdim=True) std = x.std(-1, keepdim=True) return self.gamma * (x - mean) / (std + self.eps) + self.beta class MultiHeadAttention(nn.Module): """Scaled dot-product attention over `num_heads` heads, biased by a mask.""" def __init__( self, depth: int, key_depth: int, value_depth: int, num_heads: int, bias_mask: torch.Tensor, ): super().__init__() self.num_heads = num_heads self.query_scale = (key_depth // num_heads) ** -0.5 # A buffer, not a plain attribute: it has to follow the model onto the # GPU, and it must stay out of the state dict (the checkpoint has no # entry for it) — hence persistent=False. self.register_buffer("bias_mask", bias_mask, persistent=False) self.query_linear = nn.Linear(depth, key_depth, bias=False) self.key_linear = nn.Linear(depth, key_depth, bias=False) self.value_linear = nn.Linear(depth, value_depth, bias=False) self.output_linear = nn.Linear(value_depth, depth, bias=False) def _split(self, x: torch.Tensor) -> torch.Tensor: batch, length, depth = x.shape return x.view(batch, length, self.num_heads, depth // self.num_heads).permute( 0, 2, 1, 3 ) def _merge(self, x: torch.Tensor) -> torch.Tensor: batch, heads, length, depth = x.shape return x.permute(0, 2, 1, 3).contiguous().view(batch, length, depth * heads) def forward(self, x: torch.Tensor) -> torch.Tensor: queries = self._split(self.query_linear(x)) * self.query_scale keys = self._split(self.key_linear(x)) values = self._split(self.value_linear(x)) logits = torch.matmul(queries, keys.permute(0, 1, 3, 2)) length = logits.shape[-1] logits = logits + self.bias_mask[:, :, :length, :length].type_as(logits) weights = F.softmax(logits, dim=-1) return self.output_linear(self._merge(torch.matmul(weights, values))) class Conv(nn.Module): """1-D convolution over a [batch, length, channels] sequence, centre-padded.""" def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3): super().__init__() self.pad = nn.ConstantPad1d((kernel_size // 2, (kernel_size - 1) // 2), 0) self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(self.pad(x.permute(0, 2, 1))).permute(0, 2, 1) class PositionwiseFeedForward(nn.Module): """BTC's convolutional feed-forward: conv → ReLU → conv → ReLU. The trailing ReLU is not a typo — the reference loops over its layers and activates after each one, so the block's output is non-negative. Removing it changes what the residual stream carries. """ def __init__(self, hidden_size: int, filter_size: int): super().__init__() self.layers = nn.ModuleList( [Conv(hidden_size, filter_size), Conv(filter_size, hidden_size)] ) def forward(self, x: torch.Tensor) -> torch.Tensor: for layer in self.layers: x = F.relu(layer(x)) return x class SelfAttentionBlock(nn.Module): """Pre-norm attention + feed-forward, each wrapped in a residual.""" def __init__( self, hidden_size: int, key_depth: int, value_depth: int, filter_size: int, num_heads: int, bias_mask: torch.Tensor, ): super().__init__() self.multi_head_attention = MultiHeadAttention( hidden_size, key_depth, value_depth, num_heads, bias_mask ) self.positionwise_convolution = PositionwiseFeedForward( hidden_size, filter_size ) self.layer_norm_mha = LayerNorm(hidden_size) self.layer_norm_ffn = LayerNorm(hidden_size) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x + self.multi_head_attention(self.layer_norm_mha(x)) return x + self.positionwise_convolution(self.layer_norm_ffn(x)) class BiDirectionalSelfAttention(nn.Module): """One BTC layer: a past-only and a future-only block, then a projection.""" def __init__( self, hidden_size: int, key_depth: int, value_depth: int, filter_size: int, num_heads: int, max_length: int, ): super().__init__() forward_mask = causal_mask(max_length) params = (hidden_size, key_depth, value_depth, filter_size, num_heads) self.attn_block = SelfAttentionBlock(*params, forward_mask) self.backward_attn_block = SelfAttentionBlock( *params, torch.transpose(forward_mask, 2, 3) ) self.linear = nn.Linear(hidden_size * 2, hidden_size) def forward(self, x: torch.Tensor) -> torch.Tensor: both = torch.cat((self.attn_block(x), self.backward_attn_block(x)), dim=2) return self.linear(both) class BiDirectionalSelfAttentionLayers(nn.Module): """The encoder stack: project the features in, add timing, run the layers.""" def __init__( self, feature_size: int, hidden_size: int, num_layers: int, num_heads: int, key_depth: int, value_depth: int, filter_size: int, max_length: int, ): super().__init__() self.register_buffer( "timing_signal", timing_signal(max_length, hidden_size), persistent=False ) self.embedding_proj = nn.Linear(feature_size, hidden_size, bias=False) self.self_attn_layers = nn.ModuleList( BiDirectionalSelfAttention( hidden_size, key_depth, value_depth, filter_size, num_heads, max_length, ) for _ in range(num_layers) ) self.layer_norm = LayerNorm(hidden_size) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.embedding_proj(x) x = x + self.timing_signal[:, : x.shape[1], :].type_as(x) for layer in self.self_attn_layers: x = layer(x) return self.layer_norm(x) class OutputLayer(nn.Module): """Linear projection from the encoder's hidden size to the chord vocabulary.""" def __init__(self, hidden_size: int, output_size: int): super().__init__() self.output_projection = nn.Linear(hidden_size, output_size) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.output_projection(x) # Architecture of the published checkpoints (run_config.yaml in the reference # repo). `timestep` is both the attention mask's size and the number of frames # the model is fed at a time, so it is a hard constraint on the caller, not a # tunable: muscriptor.utils.chords sizes its feature blocks around it. FEATURE_SIZE = 144 TIMESTEP = 108 HIDDEN_SIZE = 128 NUM_LAYERS = 8 NUM_HEADS = 4 FILTER_SIZE = 128 # The "large vocabulary" head: 12 roots x 14 qualities, plus X (unknown) and N. NUM_CHORDS_LARGE_VOCA = 170 class BTCModel(nn.Module): """The published BTC network, minus everything only training needed.""" def __init__(self, num_chords: int = NUM_CHORDS_LARGE_VOCA): super().__init__() self.self_attn_layers = BiDirectionalSelfAttentionLayers( FEATURE_SIZE, HIDDEN_SIZE, NUM_LAYERS, NUM_HEADS, HIDDEN_SIZE, HIDDEN_SIZE, FILTER_SIZE, TIMESTEP, ) self.output_layer = OutputLayer(HIDDEN_SIZE, num_chords) def forward(self, features: torch.Tensor) -> torch.Tensor: """Chord logits for `features` [batch, TIMESTEP, FEATURE_SIZE].""" return self.output_layer(self.self_attn_layers(features)) def load_published_state_dict(self, state_dict: dict) -> None: """Load a reference checkpoint's `model` dict into this network. The reference output layer allocates a bidirectional LSTM that its forward pass never calls, so the checkpoint carries weights this port has no module for. They are dropped here rather than loaded with `strict=False`, which would also swallow a genuinely mismatched checkpoint. """ wanted = {k: v for k, v in state_dict.items() if "output_layer.lstm." not in k} self.load_state_dict(wanted)