今年、音楽情報検索（MIR）の分野のコンペであるMIREXでコード推定タスク(Audio Chord Estimation)が久しぶりに復活しました。
面白そうなので参加してみたかったのですが、いろいろとめんどくさそうで断念しました。
久々にコード採譜モデルを作ってみたいモチベーションが高まってきたので、今回この記事で1から作成することにしました。
過去モデルを作成したときから学習データも数百件ほど増え、前回の結果を超えられるものができるはず...!!
データセット
今回、データセットは独自に採譜した約1000曲のデータを使用しました。
ジャンルはアニソン、ジャズ、ロック、ポップなどで、コード、キー、テンポなどの情報が含まれています。
前処理
ステム分離
学習時に使う音源をステム分離モデルで分離します。ボーカル、ドラム、ベース、ギター、ピアノ、その他に分割されます。
ステム分離によってモデルの精度は10%ほど上昇します。
# ステム分離モデルの読み込み
model = load_mss_model(separate_config, device=device)
for idx, wav_path in enumerate(target_files, start=1):
result = _separate_one_file(
wav_path,
output_root,
separate_config,
model,
device,
dtype,
)
ピッチシフト
ステム分離後の音源に対し、ピッチシフトを適用して学習データを水増しします（データ拡張）。ピッチシフトも精度を向上させるための重要な処理です。
y, sr = sf.read(input_path)
# pyrubberbandでピッチシフト
y_shifted = pyrb.pitch_shift(y, sr, semitone)
コードデータの正規化
元のコード表記（例: CM7, Gm, etc.）を、モデルが学習しやすい統一された形式に正規化します。これは単純にroot、bass、qualityに分割して保存するものです。
rootはCM7のCの部分、qualityはCM7のM7の部分です。
bassはCM7のようなオンコードでない場合はルートと同一で、オンコードの場合はCM7/EのEの部分です。
def convert_to_jsonl_records(spans: List[LabeledChordSpan]) -> Tuple[List[dict], int]:
"""
TSV の各行を DLchordX で解析し、JSONL レコードに変換します。
解析に失敗した行はスキップし、その件数を返します。
"""
records: List[dict] = []
skipped = 0
for span in spans:
try:
if span.chord_symbol == "N.C.":
root_text = "N"
bass_text = "N"
quality_text = "N"
else:
chord = Chord(span.chord_symbol)
root_text = chord.root.name
bass_text = chord.bass.name
quality_text = chord.quality.name
except Exception:
skipped += 1
continue
records.append(
{
"start_time": span.start_time,
"end_time": span.end_time,
"root": root_text,
"bass": bass_text,
"quality": quality_text,
}
)
return records, skipped
コードクオリティの出現頻度計算
学習時の損失関数で使用するために、データセット全体における各コードクオリティ（Major, minorなど）の出現頻度を計算します。
コードクオリティには非常に偏りがあります。例えばMajorなどのクオリティは多く出現しますが、M7(9,#11,13)などは全くありません。
これは学習時に問題になります。多く出現するクオリティばかりを出力すればロスが下がってしまうため、出現が少ないクオリティがうまく予測されません。
その偏りを改善するためにBalancedSoftmaxLossを使用します。非常にシンプルですが偏りが全く気にならないレベルまで改善されます。
class BalancedSoftmaxLoss(nn.Module):
def __init__(self, class_counts: Union[List[int], torch.Tensor], tau: float = 1.0):
"""
Args:
class_counts (Union[List[int], torch.Tensor]):
各クラスの出現回数のリストまたはテンソル。
事前に Laplace 平滑化（全カウントに+1するなど）を推奨します。
tau (float, optional): 補正のスケール係数. Defaults to 1.0.
"""
super().__init__()
class_counts = torch.tensor(class_counts, dtype=torch.float32)
# log_prior を計算し、バッファとして登録
# カウントが0のクラスは-infになるのを防ぐため、非常に小さい値にクリップ
log_prior = torch.log(torch.clamp(class_counts, min=1e-9))
self.register_buffer("log_prior", log_prior)
self.tau = tau
def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
"""
Args:
logits (torch.Tensor): モデルの出力ロジット (B, T, C)
labels (torch.Tensor): 正解ラベル (B, T, 1)
Returns:
torch.Tensor: 計算された損失値 (スカラー)
"""
# 形状を合わせる
if logits.dim() > 2:
logits = logits.reshape(-1, logits.size(-1)) # (B*T, C)
labels = labels.reshape(-1) # (B*T,)
# ロジット補正: z_k <- z_k + τ * log(n_k)
adjusted_logits = logits + self.tau * self.log_prior
loss = F.cross_entropy(adjusted_logits, labels)
return loss
フレーム単位への変換
データセットでは以下のように開始時間、終了時間の形式で記述されています。
{"start_time": 0.069, "end_time": 1.3805, "root": "N", "bass": "N", "quality": "N"}
{"start_time": 1.3805, "end_time": 1.8723, "root": "F#", "bass": "F#", "quality": ""}
{"start_time": 1.8723, "end_time": 2.528, "root": "G#", "bass": "G#", "quality": ""}
{"start_time": 2.528, "end_time": 3.1838, "root": "A", "bass": "A", "quality": "dim"}
これをモデルで学習するにはフレーム単位に変換する必要があります。
やり方として初めにframes per second(1秒当たりのフレーム数)を決定します。
仮に43だとすると、1フレーム当たりの秒数は約0.023です。
これをもとに上記のデータセットを変換すると以下のようになります。
モデルではこのフレームを予測して、正解データとのロスを計算します。
| 開始秒 (s) | 終了秒 (e) | 開始フレーム (fs) | 終了フレーム (fe) | ラベル (Root) |
|---|---|---|---|---|
| 0.069 | 1.3805 | 3 | 59 | N |
| 1.3805 | 1.8723 | 59 | 81 | F# |
| 1.8723 | 2.528 | 81 | 109 | G# |
以下は疑似コードです。
def spans_to_frames(self, spans: List[Dict[str, Any]], num_frames: int, seg_start_sec: float) -> np.ndarray:
"""
時間情報を持つイベントのリスト（spans）をフレーム系列に変換します。
Args:
spans: {"start_time": float, "end_time": float, "idx": int} のリスト
num_frames: 出力フレーム数（get_num_framesで計算された値）
seg_start_sec: オーディオセグメントの開始時間（秒）
Returns:
フレームごとのラベルインデックスを持つNumpy配列 (shape: [num_frames])
"""
# すべてのフレームをignore_indexで初期化
frame_targets = np.full((num_frames,), self.ignore_index, dtype=np.int64)
if num_frames == 0:
return frame_targets
# 各イベント（span）をループ処理
for span in spans:
event_start_sec = float(span["start_time"])
event_end_sec = float(span["end_time"])
event_idx = int(span["idx"])
# イベントの開始・終了時刻を、セグメントの開始時刻からの相対的な秒数に変換
start_relative_sec = event_start_sec - seg_start_sec
end_relative_sec = event_end_sec - seg_start_sec
# 秒数をフレームインデックスに変換
# フレームiの開始時刻は i * hop_sec なので、秒数をhop_secで割る
start_frame = math.floor(start_relative_sec / self.hop_sec)
end_frame = math.ceil(end_relative_sec / self.hop_sec)
# 計算されたフレーム範囲を、セグメントの有効範囲 [0, num_frames] にクリップ
start_frame = max(0, start_frame)
end_frame = min(num_frames, end_frame)
# 有効なフレームがあれば、その範囲のラベルをevent_idxで上書き
if end_frame > start_frame:
frame_targets[start_frame:end_frame] = event_idx
return frame_targets
モデル
Backbone（特徴抽出モデル）には、はやりの2軸方向のTransformerを採用しています。具体的には、入力テンソル $(B, T, F, C)$ に対して、$F$軸 $(B \times T, F, C)$ と $T$軸 $(B \times F, T, C)$ のそれぞれにTransformerを適用する構成です。また、Transformerのメモリ消費を抑えるため、ダウンサンプリングにより時間軸を1/8に圧縮しています。
Backboneコード
class Backbone(nn.Module):
def __init__(
self,
feature_extractor: AudioFeatureExtractor,
hidden_size: int,
output_dim: Optional[int] = None,
num_layers: int = 1,
dropout: float = 0.0,
use_gradient_checkpoint: bool = True,
):
super().__init__()
self.feature_extractor = feature_extractor
self.hidden_size = hidden_size
self.output_dim = output_dim if output_dim is not None else hidden_size
self.use_gradient_checkpoint = use_gradient_checkpoint
self.num_audio_channels = feature_extractor.num_audio_channels
self.n_bins = feature_extractor.n_bins
self.oct_frontend = OctaveSharedAggregate(
in_channels=self.num_audio_channels, # 12
f_total=self.n_bins, # 252
hidden_size=hidden_size,
bins_per_octave=feature_extractor.bins_per_octave,
octave_emb_dim=32,
dropout=dropout,
use_film=True,
)
self.conv1 = nn.Conv2d(hidden_size, hidden_size, kernel_size=3, padding=1)
self.down_conv = nn.Sequential(
nn.Conv2d(hidden_size, hidden_size * 2, kernel_size=3, padding=1, stride=(2, 1)),
nn.GroupNorm(4, hidden_size * 2),
nn.GELU(),
nn.Conv2d(
hidden_size * 2,
hidden_size * 4,
kernel_size=3,
padding=1,
stride=(2, 1),
),
nn.GroupNorm(4, hidden_size * 4),
nn.GELU(),
nn.Conv2d(
hidden_size * 4,
hidden_size * 4,
kernel_size=3,
padding=1,
stride=(2, 1),
),
nn.GroupNorm(4, hidden_size * 4),
nn.GELU(),
nn.Conv2d(hidden_size * 4, hidden_size * 4, kernel_size=3, padding=1),
nn.GroupNorm(4, hidden_size * 4),
)
self.layers = nn.ModuleList([])
for _ in range(num_layers):
time_roformer = Transformer(
input_dim=hidden_size * 4,
head_dim=hidden_size * 4 // 8,
num_layers=1,
num_heads=8,
ffn_hidden_size_factor=2,
dropout=dropout,
)
band_roformer = Transformer(
input_dim=hidden_size * 4,
head_dim=hidden_size * 4 // 8,
num_layers=1,
num_heads=8,
ffn_hidden_size_factor=2,
dropout=dropout,
)
self.layers.append(nn.ModuleList([time_roformer, band_roformer]))
self.final_norm = RMSNorm(hidden_size * 4)
freq_downsample_factor = 1
input_freq_bins = feature_extractor.bins_per_octave
output_freq_bins = input_freq_bins // freq_downsample_factor
final_channels = hidden_size * 4
flattened_dim = output_freq_bins * final_channels
self.to_time_features = nn.Conv1d(flattened_dim, self.output_dim, kernel_size=1)
self.up_time = nn.ConvTranspose1d(
self.output_dim,
self.output_dim,
kernel_size=8,
stride=8,
)
def forward(
self,
waveform: torch.Tensor,
context: Optional[BackboneContext] = None,
return_intermediate: bool = False,
) -> Union[torch.Tensor, Tuple[torch.Tensor, List[torch.Tensor]]]:
if context is None:
context = self.feature_extractor(waveform)
use_checkpoint = self.use_gradient_checkpoint and self.training and torch.is_grad_enabled()
checkpoint_fn = torch.utils.checkpoint.checkpoint if use_checkpoint else checkpoint_bypass
x = self.oct_frontend(context.spec) # (B, hidden_size, T, 36)
x = self.conv1(x)
# ダウンサンプリング
x = self.down_conv(x)
x = einops.rearrange(x, "b c t f -> b t f c") # (B, downT, F, C)
intermediate_features = []
for time_roformer, band_roformer in self.layers:
B, T, freq, C = x.shape
# 周波数軸Transformer
x = x.reshape(B * T, freq, C) # [B*T, F, C]
x = checkpoint_fn(band_roformer, x, use_reentrant=False)
x = x.reshape(B, T, freq, C)
# 時間軸Transformer
x = einops.rearrange(x, "b t f c -> (b f) t c") # [B*F, T, C]
x = checkpoint_fn(time_roformer, x, use_reentrant=False)
x = einops.rearrange(x, "(b f) t c -> b t f c", f=freq) # [B, T, F, C]
if return_intermediate:
intermediate_features.append(x)
x = self.final_norm(x) # [B, T, F, D]
x = einops.rearrange(x, "b t f d -> b (f d) t") # [B, F*D, downT]
x = self.to_time_features(x) # [B, output_dim, downT]
x = self.up_time(x) # [B, output_dim, downT*8]
target_T = context.crop_length
if x.shape[-1] < target_T:
x = F.pad(x, (0, target_T - x.shape[-1]))
else:
x = x[..., :target_T]
x = einops.rearrange(x, "b d t -> b t d") # [B, T, output_dim]
if return_intermediate:
return x, intermediate_features
return x
BaseTranscriptionModelは、ベースとなる採譜モデルです。 主な役割は、chord25を介してroot_chordとbassを予測することと、Refinerによる平滑化です。
chord25はコードのピッチクラスとベースを合わせた表現で、データの偏りが起きにくいという特徴があります。この表現をベースに予測を行うことで精度を安定させています。 また、Refinerでchord25の出力を滑らかにしています。これがないと、後段でCRFを適用した際に十分な平滑化効果が得られません。
学習は2段階で構成されています。
第1段階: Backbone + BaseTranscriptionModel で学習
第2段階: 学習済みBackboneを読み込み、Backbone + SegmentTranscriptionModel を学習
BaseTranscriptionModelコード
class BaseTranscriptionModel(nn.Module):
def __init__(
self,
backbone: Backbone,
hidden_size: int,
num_quality_classes: int,
num_root_classes: int,
num_bass_classes: int,
num_key_classes: int,
num_tempo_classes: Optional[int] = None,
dropout_probability: float = 0.0,
use_layer_norm: bool = True,
) -> None:
super().__init__()
self.backbone = backbone
self.hidden_size = hidden_size
self.num_quality_classes = num_quality_classes
self.num_root_classes = num_root_classes
self.norm = RMSNorm(hidden_size) if use_layer_norm else nn.Identity()
self.dropout = nn.Dropout(dropout_probability) if dropout_probability > 0.0 else nn.Identity()
self.num_root_quality_classes = (num_quality_classes - 1) * 12 + 1
self.num_bass_classes = num_bass_classes
# Main Branch Heads (backbone特徴量から予測)
chord25_hidden = 25 + 12
self.chord25_head = nn.Sequential(nn.Linear(hidden_size, chord25_hidden))
self.boundary_head = nn.Linear(hidden_size, 1)
self.key_head = nn.Linear(hidden_size, num_key_classes)
if num_tempo_classes is None:
self.tempo_head = nn.Linear(hidden_size, 1)
self.tempo_is_regression = True
else:
self.tempo_head = nn.Linear(hidden_size, num_tempo_classes)
self.tempo_is_regression = False
# --- Refiner Branch ---
smooth_hidden = 256
self.smooth_stage1 = nn.Sequential(
nn.Linear(chord25_hidden, smooth_hidden),
RMSNorm(smooth_hidden),
nn.Tanh(), # stage_1_activation
)
self.smooth_rnn = nn.LSTM(smooth_hidden, smooth_hidden // 2, num_layers=1, batch_first=True, bidirectional=True)
self.smooth_net_proj = nn.Sequential(
nn.Linear(smooth_hidden * 2, smooth_hidden),
RMSNorm(smooth_hidden),
nn.ReLU(inplace=True),
)
self.smooth_head = nn.Sequential(nn.Linear(smooth_hidden, chord25_hidden))
self.root_chord_head = nn.Linear(chord25_hidden, self.num_root_quality_classes)
self.bass_head = nn.Linear(chord25_hidden, num_bass_classes)
def _compute_heads_and_answer(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
normed = self.norm(features) # RMSNorm or Identity
features = self.dropout(normed)
# Primary heads
chord25_logits = self.chord25_head(features)
chord25_logits = torch.clamp(torch.relu(chord25_logits), max=1.0)
outputs = {
"initial_chord25_logits": chord25_logits[..., :25],
"initial_chord25_original": chord25_logits,
"initial_key_logits": self.key_head(features),
"initial_tempo": self.tempo_head(features),
"initial_boundary_logits": self.boundary_head(features),
}
# Smooth Branch
# chord25_logits: (B, T, 25)
# 1. Stage 1 (Dense Tanh)
# Linear: (B, T, 25) -> (B, T, 256)
x = self.smooth_stage1(chord25_logits)
# 2. RNN
skip = x
self.smooth_rnn.flatten_parameters()
x, _ = self.smooth_rnn(x) # (B, T, 256)
x = torch.concat([skip, x], dim=-1)
# 3. Stage 2 (Dense ReLU)
# Linear: (B, T, 256) -> (B, T, 256)
x = self.smooth_net_proj(x)
# 4. Smooth Head (Dense BN ReLU 25)
# Linear: (B, T, 256) -> (B, T, 25)
smooth_features = self.smooth_head(x)
smooth_features = torch.clamp(torch.relu(smooth_features), max=1.0)
# Root Chord from Smooth Features
outputs["initial_root_chord_logits"] = self.root_chord_head(smooth_features)
outputs["initial_bass_logits"] = self.bass_head(smooth_features)
# Smooth Chord25 Logits (for loss calculation)
outputs["initial_smooth_chord25_logits"] = smooth_features[..., :25]
outputs["initial_smooth_chord25_original"] = smooth_features
return outputs
def forward(
self, waveform: torch.Tensor, global_step: Optional[int] = None, max_segments: int = 256
) -> Dict[str, torch.Tensor]:
# Backbone Feature Extraction
features = self.backbone(waveform)
# Heads & Segment Logic
outputs = self._compute_heads_and_answer(features)
return outputs
SegmentTranscriptionModelの役割は、予測結果に文脈を反映させることです。
文脈に強いTransformerを適用するにあたり、2つの課題がありました。 1つはメモリ使用量です。Backboneの特徴量をそのまま入力すると、系列長が長くなりすぎてメモリを圧迫します。 もう1つは、特徴量がフレーム単位であるためノイズの影響を受けやすく、コード間の関係性を必ずしも効率的に学べるとは限らない点です。
そこで、コード境界（chord_boundary）で特徴量をセグメント分割し、区間内を平均化する処理を加えました。これにより、1コード1フレームの系列として集約しています。 このアプローチによってメモリ使用量を削減しつつ、コード間の関係性をより明確に学習できる構造にしました。
学習したattentionの重みを見てみると、少し先の領域を見たり、少し前の領域を見たり、連続でコードを参照したりと面白いです。
SegmentTranscriptionModelコード
import torch
import torch.nn.functional as F
import torch.nn as nn
import torch.utils.checkpoint
from typing import Optional, Dict, Tuple, List
from .transformer import RMSNorm, Transformer
from .transcription_model import Backbone
def checkpoint_bypass(func, *args, **kwargs):
"""チェックポイントを使用しない場合のバイパス関数"""
return func(*args)
class SegmentFeatureProcessor:
"""
セグメントに基づいた特徴量の集約（プーリング）、パディング、復元（ブロードキャスト）を行うクラス。
計算ロジック（どう計算するか）のみに責務を持つ。
"""
def aggregate_features(
self,
frame_features: torch.Tensor,
segment_ids: torch.Tensor,
num_segments: int,
epsilon: float = 1e-8,
) -> torch.Tensor:
"""
1サンプル分の特徴量をセグメント単位で平均プーリングする。
Args:
frame_features: (TotalFrames, FeatureDim)
segment_ids: (TotalFrames,) 各フレームの所属セグメントID
num_segments: セグメント総数
Returns:
aggregated_features: (NumSegments, FeatureDim)
"""
# Tensorの形状とデバイス情報を取得
total_frames, feature_dim = frame_features.shape
device = frame_features.device
dtype = frame_features.dtype
# 集計用バッファの初期化
sum_features = torch.zeros((num_segments, feature_dim), device=device, dtype=dtype)
frame_counts = torch.zeros((num_segments, 1), device=device, dtype=dtype)
# IDに基づいて特徴量を加算 (Scatter Add)
sum_features.index_add_(0, segment_ids, frame_features)
# IDごとのフレーム数をカウント
ones = torch.ones((total_frames, 1), device=device, dtype=dtype)
frame_counts.index_add_(0, segment_ids, ones)
# 平均を計算 (ゼロ除算防止)
return sum_features / (frame_counts + epsilon)
def collate_segments(
self,
segment_feature_list: List[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
可変長のセグメント特徴量リストをパディングしてバッチ化する。
Returns:
padded_features: (BatchSize, MaxSegments, FeatureDim)
padding_mask: (BatchSize, MaxSegments) 有効な場所がTrue
"""
batch_size = len(segment_feature_list)
if batch_size == 0:
return torch.empty(0), torch.empty(0)
feature_dim = segment_feature_list[0].shape[-1]
# バッチ内で最大のセグメント数を取得
max_segments = max(feat.shape[0] for feat in segment_feature_list)
device = segment_feature_list[0].device
dtype = segment_feature_list[0].dtype
padded_features = torch.zeros((batch_size, max_segments, feature_dim), device=device, dtype=dtype)
padding_mask = torch.zeros((batch_size, max_segments), device=device, dtype=torch.bool)
for i, features in enumerate(segment_feature_list):
num_seg = features.shape[0]
padded_features[i, :num_seg] = features
padding_mask[i, :num_seg] = True
return padded_features, padding_mask
def broadcast_to_frames(
self,
padded_segment_features: torch.Tensor,
segment_ids_batch: torch.Tensor,
) -> torch.Tensor:
"""
セグメント特徴量をフレーム単位に引き伸ばす（アップサンプリング）。
Args:
padded_segment_features: (BatchSize, MaxSegments, FeatureDim)
segment_ids_batch: (BatchSize, TotalFrames) フレームごとの所属ID
Returns:
frame_features: (BatchSize, TotalFrames, FeatureDim)
"""
batch_size, total_frames = segment_ids_batch.shape
feature_dim = padded_segment_features.shape[-1]
output_frame_features = torch.zeros(
(batch_size, total_frames, feature_dim),
device=padded_segment_features.device,
dtype=padded_segment_features.dtype,
)
# バッチごとのループ処理（可読性と安全性を優先）
# ※ gatherを使ったベクトル化も可能だが、メモリ消費と可読性のトレードオフを考慮しループを採用
for i in range(batch_size):
# (TotalFrames,)
ids = segment_ids_batch[i]
# IDに対応する特徴量をコピー: (TotalFrames, FeatureDim)
output_frame_features[i] = padded_segment_features[i].index_select(0, ids)
return output_frame_features
class BatchBoundarySegmenter:
"""
バッチデータに対する「境界検出」から「特徴量プーリング」までを一括管理するクラス。
"""
def __init__(
self,
threshold: float = 0.5,
nms_window_radius: int = 3,
min_segment_length: int = 8,
max_segments: Optional[int] = None,
feature_processor: Optional[SegmentFeatureProcessor] = None,
):
# 境界検出のパラメータ
self.threshold = threshold
self.nms_window_radius = nms_window_radius
self.min_segment_length = min_segment_length
self.max_segments = max_segments
# 計算ロジック担当のクラス（DI: Dependency Injection）
self.processor = feature_processor or SegmentFeatureProcessor()
def set_max_segments(self, max_segments: int):
self.max_segments = max_segments
@torch.no_grad()
def process_batch(
self,
frame_features: torch.Tensor, # (B, T, D)
boundary_logits: torch.Tensor, # (B, T) or (B, T, 1)
detach_boundary: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[List[Tuple[int, int]]]]:
"""
バッチ処理の実行メソッド。
Returns:
padded_features: (B, S_max, D) セグメントごとの特徴量（パディング済み）
padding_mask: (B, S_max) マスク
segment_ids_batch: (B, T) 各フレームのセグメントID
segments_info_batch: List[List[(start, end)]] 区間情報のリスト
"""
# 入力形状の正規化
if boundary_logits.dim() == 3:
boundary_logits = boundary_logits.squeeze(-1)
batch_size, total_frames, feature_dim = frame_features.shape
# 確率値への変換
probabilities = torch.sigmoid(boundary_logits)
if detach_boundary:
probabilities = probabilities.detach()
# 結果格納用
segment_features_list: List[torch.Tensor] = []
segment_ids_list: List[torch.Tensor] = []
segments_info_batch: List[List[Tuple[int, int]]] = []
for i in range(batch_size):
# 1. 境界検出
boundary_indices = self._detect_boundaries_1d(probabilities[i])
# 2. セグメント情報への変換
segments, seg_ids = self._convert_to_segments(boundary_indices, total_frames)
# デバイス転送
seg_ids = seg_ids.to(frame_features.device)
num_segments = len(segments)
# 3. 特徴量プーリング (Processorに委譲)
pooled_feat = self.processor.aggregate_features(frame_features[i], seg_ids, num_segments)
segment_features_list.append(pooled_feat)
segment_ids_list.append(seg_ids)
segments_info_batch.append(segments)
# 4. パディング処理 (Processorに委譲)
padded_features, padding_mask = self.processor.collate_segments(segment_features_list)
# 5. IDリストをTensorにスタック (フレーム数固定の前提により可能)
segment_ids_batch = torch.stack(segment_ids_list)
return padded_features, padding_mask, segment_ids_batch, segments_info_batch
def expand_to_frames(self, padded_segment_features: torch.Tensor, segment_ids_batch: torch.Tensor) -> torch.Tensor:
"""
セグメント特徴量をフレーム単位に復元する。
実際の処理は internal processor に委譲する。
"""
return self.processor.broadcast_to_frames(padded_segment_features, segment_ids_batch)
# --- 内部メソッド: 境界検出ロジック ---
def _detect_boundaries_1d(self, probabilities: torch.Tensor) -> List[int]:
"""1サンプル分の境界検出（NMS + Greedy）"""
sequence_length = probabilities.numel()
if sequence_length <= 1:
return []
# NMSによるピーク抽出
win_size = 2 * self.nms_window_radius + 1
max_pooled = F.max_pool1d(
probabilities.view(1, 1, sequence_length), kernel_size=win_size, stride=1, padding=self.nms_window_radius
).view(sequence_length)
is_peak = (probabilities >= max_pooled) & (probabilities > self.threshold)
candidates = torch.nonzero(is_peak).flatten()
candidates = candidates[(candidates > 0) & (candidates < sequence_length)] # 始点終点除外
if candidates.numel() == 0:
return []
# 距離制約によるフィルタリング
# スコア順にソート
scores = probabilities[candidates]
sorted_indices = candidates[torch.argsort(scores, descending=True)].tolist()
selected: List[int] = []
for t in sorted_indices:
if self.max_segments and len(selected) >= (self.max_segments - 1):
break
# 制約チェック
if t < self.min_segment_length:
continue
if (sequence_length - t) < self.min_segment_length:
continue
is_far_enough = True
for s in selected:
if abs(t - s) < self.min_segment_length:
is_far_enough = False
break
if is_far_enough:
selected.append(t)
selected.sort()
return selected
def _convert_to_segments(
self, boundary_indices: List[int], total_frames: int
) -> Tuple[List[Tuple[int, int]], torch.Tensor]:
"""境界位置リストからセグメント区間とIDマップを生成"""
cut_points = [0] + boundary_indices + [total_frames]
segments = []
for i in range(len(cut_points) - 1):
s, e = int(cut_points[i]), int(cut_points[i + 1])
if e > s:
segments.append((s, e))
seg_ids = torch.empty((total_frames,), dtype=torch.long)
for idx, (s, e) in enumerate(segments):
seg_ids[s:e] = idx
return segments, seg_ids
class SegmentTranscriptionModel(nn.Module):
def __init__(
self,
backbone: Backbone,
hidden_size: int,
num_quality_classes: int,
num_root_classes: int,
num_bass_classes: int,
num_key_classes: int,
num_tempo_classes: Optional[int] = None,
dropout_probability: float = 0.0,
use_layer_norm: bool = True,
transformer_hidden_size: int = 256,
transformer_num_heads: int = 8,
transformer_num_layers: int = 3,
) -> None:
super().__init__()
self.backbone = backbone
self.hidden_size = hidden_size
self.num_quality_classes = num_quality_classes
self.num_root_classes = num_root_classes
self.norm = RMSNorm(hidden_size) if use_layer_norm else nn.Identity()
self.dropout = nn.Dropout(dropout_probability) if dropout_probability > 0.0 else nn.Identity()
self.num_root_quality_classes = (num_quality_classes - 1) * 12 + 1
self.num_bass_classes = num_bass_classes
# Main Branch Heads (backbone特徴量から予測)
self.boundary_head = nn.Linear(hidden_size, 1)
if num_tempo_classes is None:
self.tempo_head = nn.Linear(hidden_size, 1)
self.tempo_is_regression = True
else:
self.tempo_head = nn.Linear(hidden_size, num_tempo_classes)
self.tempo_is_regression = False
# --- Segment Branch ---
self.segmenter = BatchBoundarySegmenter(
threshold=0.5, nms_window_radius=3, min_segment_length=8, max_segments=256
)
self.seg_transformer = Transformer(
input_dim=hidden_size,
head_dim=transformer_hidden_size // transformer_num_heads,
num_heads=transformer_num_heads,
num_layers=transformer_num_layers,
ffn_hidden_size_factor=2,
)
self.seg_out_proj = nn.Linear(hidden_size, hidden_size)
self.key_head = nn.Linear(hidden_size, num_key_classes)
self.smooth_head = nn.Sequential(nn.Linear(hidden_size, 25 + 12))
self.root_chord_head = nn.Linear(25 + 12, self.num_root_quality_classes)
self.bass_head = nn.Linear(25 + 12, num_bass_classes)
def _ste_broadcast(self, padded_features, seg_ids, boundary_logits):
# 1. 形状調整
if boundary_logits.dim() == 3:
boundary_logits = boundary_logits.squeeze(-1)
max_segments = padded_features.shape[1]
device = padded_features.device
# 2. 位置情報の準備
probs = torch.sigmoid(boundary_logits)
# pos_raw: 勾配の源 (float)
pos_raw = torch.cumsum(probs, dim=1) # (B, T)
# pos_anchor: 値の固定先 (int -> float)
pos_anchor = seg_ids.float() # (B, T)
# 値は pos_anchor (整数)、勾配は pos_raw (確率) に流す
pos_aligned = pos_anchor + (pos_raw - pos_raw.detach())
# 3. Soft重みの計算 (Gaussian / Softmax Kernel)
# (B, T, 1)
current_pos = pos_aligned.unsqueeze(-1)
# (1, 1, S)
target_idx = torch.arange(max_segments, device=device).view(1, 1, -1)
# 距離の二乗 (B, T, S)
dist_sq = (current_pos - target_idx) ** 2
# sigma相当の係数を掛けて鋭さを調整 (大きいほど鋭い)
sharpness = 10.0
weight_soft = F.softmax(-sharpness * dist_sq, dim=-1)
# 4. Hard重みの計算 (Forward用)
weight_hard = F.one_hot(seg_ids, num_classes=max_segments).float()
# 5. 重みレベルでのSTE
weight_ste = weight_hard + (weight_soft - weight_soft.detach())
# 6. 復元
return torch.matmul(weight_ste, padded_features)
def _compute_heads_and_answer(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
use_checkpoint = self.training and torch.is_grad_enabled()
checkpoint_fn = torch.utils.checkpoint.checkpoint if use_checkpoint else checkpoint_bypass
normed = self.norm(features) # RMSNorm or Identity
features = self.dropout(normed)
boundary_logits = self.boundary_head(features)
outputs = {
"initial_tempo": self.tempo_head(features),
"initial_boundary_logits": boundary_logits,
}
seg_pad, seg_mask, seg_ids, segments = self.segmenter.process_batch(
frame_features=features, boundary_logits=boundary_logits, detach_boundary=True
)
x = seg_pad * seg_mask.unsqueeze(-1).to(seg_pad.dtype)
x = checkpoint_fn(self.seg_transformer, x, use_reentrant=False)
x = self.seg_out_proj(x)
if self.training and torch.is_grad_enabled():
frame_ctx = self._ste_broadcast(x, seg_ids, boundary_logits)
else:
frame_ctx = self.segmenter.expand_to_frames(x, seg_ids)
refined_features = features + frame_ctx
outputs["initial_key_logits"] = self.key_head(refined_features)
smooth_features = self.smooth_head(refined_features)
smooth_features = torch.clamp(torch.relu(smooth_features), max=1.0)
# Root Chord from Smooth Features
outputs["initial_root_chord_logits"] = self.root_chord_head(smooth_features)
outputs["initial_bass_logits"] = self.bass_head(smooth_features)
# Smooth Chord25 Logits (for loss calculation)
outputs["initial_smooth_chord25_logits"] = smooth_features[..., :25]
outputs["initial_smooth_chord25_original"] = smooth_features
return outputs
def forward(
self, waveform: torch.Tensor, global_step: Optional[int] = None, max_segments: int = 256
) -> Dict[str, torch.Tensor]:
self.segmenter.set_max_segments(max_segments)
# Backbone Feature Extraction
features = self.backbone(waveform)
# Heads & Segment Logic
outputs = self._compute_heads_and_answer(features)
return outputs
手法追加による精度変化（改善ポイント）
ベースラインからどのような手法で精度が向上したかをまとめました。厳密な比較はめんどくさいので行っていませんが大体の目安としてみてください。
| Task | Baseline | + ステム拡張, SpecAugment | + CQT入力 | + Refiner | + ピッチシフト | + 入力長拡大 (最終) |
|---|---|---|---|---|---|---|
| bass | 69.16 | 78.21 (+9.05) | 88.94 (+10.73) | 90.00 (+1.06) | 91.78 (+1.78) | 92.79 (+1.01) |
| root | 70.39 | 77.55 (+7.16) | 86.48 (+8.93) | 87.73 (+1.25) | 89.74 (+2.01) | 90.60 (+0.86) |
| quality | 53.85 | 59.01 (+5.16) | 67.64 (+8.63) | 68.25 (+0.61) | 70.73 (+2.48) | 73.04 (+2.31) |
| key | 75.81 | 74.03 (-1.78) | 80.40 (+6.37) | 79.87 (-0.53) | 82.35 (+2.48) | 88.84 (+6.49) |
ステム拡張、SpecAugment
ステム拡張は、6つのステムをミックスする際に音量をランダムで変更したり、特定のステムをドロップするものです。これによって特定の楽器が無い場合の音源が再現でき、モデルの文脈予測を改善できるはずです。例えば、ベースのステムがドロップされた場合、モデルはベース音の推定にベースステムを使うことができず、ほかの楽器の音や文脈から推測する必要があります。
# 必要なステムのみ読み込む（ステムミュート）
stems, stems_to_mute = ensure_required_chord_stem(
stems=stems,
stem_order=self.stem_order,
required_indices=self.required_stem_indices,
stems_to_mute=stems_to_mute,
stems_dir=stems_dir,
target_sample_rate=self.sample_rate,
start_seconds=segment_start_sec,
segment_seconds=self.segment_seconds,
shift_semitones=shift_semitones,
)
# ランダム・ゲイン
if random.random() < self.mixing_cfg.get("p", 0.0):
min_db = self.mixing_cfg.get("min_gain_db", -6.0)
max_db = self.mixing_cfg.get("max_gain_db", 3.0)
for i in self.nondrum_indices:
gain_db = random.uniform(min_db, max_db)
gain_linear = 10.0 ** (gain_db / 20.0)
stems[i] *= gain_linear
SpecAugmentはランダムにスペクトログラムの周波数方向、時間軸方向をマスクします。これはステムごとに行われます。ステム拡張と同様に欠落した情報を周りから推測する必要があり、モデルの文脈予測改善につながると思います。
batch_size = waveform.shape[0]
stem_specs: List[torch.Tensor] = []
for stem_idx in range(self.num_stems):
stem_waveform = waveform_per_stem[:, stem_idx] # [B, C_per_stem, T]
stem_flat = einops.rearrange(stem_waveform, "b c t -> (b c) t")
stem_cqt = self.cqt(stem_flat.float(), return_complex=False)
stem_cqt = einops.rearrange(
stem_cqt, "(b c) f t -> b c f t", b=batch_size, c=self.channels_per_stem
).contiguous()
# SpecAugment
stem_cqt_aug, _ = self.spec_augment(stem_cqt)
stem_specs.append(stem_cqt_aug)
spec = torch.cat(stem_specs, dim=1)
また、この手法を追加した際に面白い現象が発生しました。
通常データ拡張をするとトレーニングロスは高くなることが多いですが、今回のケースではトレーニングロスが低くなりました。
データの多様性が上がってトレーニングが安定したというのもあると思いますが、文脈予測を強化するような学習が行われるからというのもありそうです。
CQT入力
特に改善が大きかったのが、入力をMel SpectrogramからCQT Spectrogramに変更することでした。
MelSpectrogramに比べてCQTは周波数方向の分解能が高いので、よりピッチを推定しやすいということだと思います。コード推定タスクではこのCQTが良く使われるのも納得です。特に理由がない限りCQTを使うべきですね。再帰ダウンサンプリングCQTであればMel Spectrogramを計算するのと計算コストも大きく変わりません。
Refiner
Refinerはモデルの推論を2段階に分けます。1段目は通常通り予測し、2段目は1段目のモデルの最終出力(chord25)を入力とし、全体を見て修正を行う役割を担います。
指標上は改善が小さいですが、より滑らかな出力を得られます。
ピッチシフト
ピッチシフトはオーディオを-5~+6までピッチシフトしてデータを拡張します。1つのキーのデータで12種類のキーのデータが作成できるので、そのコードが持つ、キーに対しての相対的な意味(キーが変わってもそのキーに対する機能は変わらない)を学習するのを助けるはず...です。適用コストは高いので事前にオーディオに対して計算しておきます。
入力長拡大
入力長拡大では入力の長さを8秒から60秒に拡大します。この手法ではキーの指標が大きく改善しました。キーの予測には長い文脈が必要だとわかります。
例えばですが、コードの遷移の仕方や、メロディの使い方など、短い区間だけでは考慮しづらいということだと思います。(トニック化や借用など短い区間では判定しずらいケースがある)
おわりに
久しぶりにコード採譜モデルを学習させてみました。前回の時と比べて知見が広がったためより精度の高いものが作成できました。
さらに今回は学習コードも公開したのでデータセットがある方は自分でも学習できると思います。
作成したコードはGithubで公開しています。