# Copyright 2024 The YourMT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Please see the details in the LICENSE file.
"""ymt3.py"""
import os
from typing import Union, Optional, Tuple, Dict, List, Any
from collections import Counter
import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
import torchaudio # for debugging audio
import pytorch_lightning as pl
import numpy as np
import wandb
from einops import rearrange
from transformers import T5Config
from model.t5mod import T5EncoderYMT3, T5DecoderYMT3, MultiChannelT5Decoder
from model.t5mod_helper import task_cond_dec_generate
from model.perceiver_mod import PerceiverTFEncoder
from model.perceiver_helper import PerceiverTFConfig
from model.conformer_mod import ConformerYMT3Encoder
from model.conformer_helper import ConformerYMT3Config
from model.lm_head import LMHead
from model.pitchshift_layer import PitchShiftLayer
from model.spectrogram import get_spectrogram_layer_from_audio_cfg
from model.conv_block import PreEncoderBlockRes3B
from model.conv_block import PreEncoderBlockHFTT, PreEncoderBlockRes3BHFTT # added for hFTT-like pre-encoder
from model.projection_layer import get_projection_layer, get_multi_channel_projection_layer
from model.optimizers import get_optimizer
from model.lr_scheduler import get_lr_scheduler
from utils.note_event_dataclasses import Note
from utils.note2event import mix_notes
from utils.event2note import merge_zipped_note_events_and_ties_to_notes, DECODING_ERR_TYPES
from utils.metrics import compute_track_metrics
from utils.metrics import AMTMetrics
# from utils.utils import write_model_output_as_npy
from utils.utils import write_model_output_as_midi, create_inverse_vocab, write_err_cnt_as_json
from utils.utils import Timer
from utils.task_manager import TaskManager
from config.config import audio_cfg as default_audio_cfg
from config.config import model_cfg as default_model_cfg
from config.config import shared_cfg as default_shared_cfg
from config.config import T5_BASE_CFG
class YourMT3(pl.LightningModule):
"""YourMT3:
Lightning wrapper for multi-task music transcription Transformer.
"""
def __init__(
self,
audio_cfg: Optional[Dict] = None,
model_cfg: Optional[Dict] = None,
shared_cfg: Optional[Dict] = None,
pretrained: bool = False,
optimizer_name: str = 'adamwscale',
scheduler_name: str = 'cosine',
base_lr: float = None, # None: 'auto' for AdaFactor, 1e-3 for constant, 1e-2 for cosine
max_steps: Optional[int] = None,
weight_decay: float = 0.0,
init_factor: Optional[Union[str, float]] = None,
task_manager: TaskManager = TaskManager(),
eval_subtask_key: Optional[str] = "default",
eval_vocab: Optional[Dict] = None,
eval_drum_vocab: Optional[Dict] = None,
write_output_dir: Optional[str] = None,
write_output_vocab: Optional[Dict] = None,
onset_tolerance: float = 0.05,
add_pitch_class_metric: Optional[List[str]] = None,
add_melody_metric_to_singing: bool = True,
test_optimal_octave_shift: bool = False,
test_pitch_shift_layer: Optional[str] = None,
**kwargs: Any) -> None:
super().__init__()
if pretrained is True:
raise NotImplementedError("Pretrained model is not supported in this version.")
self.test_pitch_shift_layer = test_pitch_shift_layer # debug only
# Config
if model_cfg is None:
model_cfg = default_model_cfg # default config, not overwritten by args of trainer
if audio_cfg is None:
audio_cfg = default_audio_cfg # default config, not overwritten by args of trainer
if shared_cfg is None:
shared_cfg = default_shared_cfg # default config, not overwritten by args of trainer
# Spec Layer (need to define here to infer max token length)
self.spectrogram, spec_output_shape = get_spectrogram_layer_from_audio_cfg(
audio_cfg) # can be spec or melspec; output_shape is (T, F)
model_cfg["feat_length"] = spec_output_shape[0] # T of (T, F)
# Task manger and Tokens
self.task_manager = task_manager
self.max_total_token_length = self.task_manager.max_total_token_length
# Task Conditioning
self.use_task_cond_encoder = bool(model_cfg["use_task_conditional_encoder"])
self.use_task_cond_decoder = bool(model_cfg["use_task_conditional_decoder"])
# Select Encoder type, Model-specific Config
assert model_cfg["encoder_type"] in ["t5", "perceiver-tf", "conformer"]
assert model_cfg["decoder_type"] in ["t5", "multi-t5"]
self.encoder_type = model_cfg["encoder_type"] # {"t5", "perceiver-tf", "conformer"}
self.decoder_type = model_cfg["decoder_type"] # {"t5", "multi-t5"}
encoder_config = model_cfg["encoder"][self.encoder_type] # mutable
decoder_config = model_cfg["decoder"][self.decoder_type] # mutable
# Positional Encoding
if isinstance(model_cfg["num_max_positions"], str) and model_cfg["num_max_positions"] == 'auto':
encoder_config["num_max_positions"] = int(model_cfg["feat_length"] +
self.task_manager.max_task_token_length + 10)
decoder_config["num_max_positions"] = int(self.max_total_token_length + 10)
else:
assert isinstance(model_cfg["num_max_positions"], int)
encoder_config["num_max_positions"] = model_cfg["num_max_positions"]
decoder_config["num_max_positions"] = model_cfg["num_max_positions"]
# Select Pre-Encoder and Pre-Decoder type
if model_cfg["pre_encoder_type"] == "default":
model_cfg["pre_encoder_type"] = model_cfg["pre_encoder_type_default"].get(model_cfg["encoder_type"], None)
elif model_cfg["pre_encoder_type"] in [None, "none", "None", "0"]:
model_cfg["pre_encoder_type"] = None
if model_cfg["pre_decoder_type"] == "default":
model_cfg["pre_decoder_type"] = model_cfg["pre_decoder_type_default"].get(model_cfg["encoder_type"]).get(
model_cfg["decoder_type"], None)
elif model_cfg["pre_decoder_type"] in [None, "none", "None", "0"]:
model_cfg["pre_decoder_type"] = None
self.pre_encoder_type = model_cfg["pre_encoder_type"]
self.pre_decoder_type = model_cfg["pre_decoder_type"]
# Pre-encoder
self.pre_encoder = nn.Sequential()
if self.pre_encoder_type in ["conv", "conv1d_t", "conv1d_f"]:
kernel_size = (3, 3)
avp_kernel_size = (1, 2)
if self.pre_encoder_type == "conv1d_t":
kernel_size = (3, 1)
elif self.pre_encoder_type == "conv1d_f":
kernel_size = (1, 3)
self.pre_encoder.append(
PreEncoderBlockRes3B(1,
model_cfg["conv_out_channels"],
kernel_size=kernel_size,
avp_kernerl_size=avp_kernel_size,
activation="relu"))
pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1] // 2**3, model_cfg["conv_out_channels"]
) # (T, F, C) excluding batch dim
elif self.pre_encoder_type == "hftt":
self.pre_encoder.append(PreEncoderBlockHFTT())
pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1], 128) # (T, F, C) excluding batch dim
elif self.pre_encoder_type == "res3b_hftt":
self.pre_encoder.append(PreEncoderBlockRes3BHFTT())
pre_enc_output_shape = (spec_output_shape[0], spec_output_shape[1] // 2**3, 128)
else:
pre_enc_output_shape = spec_output_shape # (T, F) excluding batch dim
# Auto-infer `d_feat` and `d_model`, `vocab_size`, and `num_max_positions`
if isinstance(model_cfg["d_feat"], str) and model_cfg["d_feat"] == 'auto':
if self.encoder_type == "perceiver-tf" and encoder_config["attention_to_channel"] is True:
model_cfg["d_feat"] = pre_enc_output_shape[-2] # TODO: better readablity
else:
model_cfg["d_feat"] = pre_enc_output_shape[-1] # C of (T, F, C) or F or (T, F)
if self.encoder_type == "perceiver-tf" and isinstance(encoder_config["d_model"], str):
if encoder_config["d_model"] == 'q':
encoder_config["d_model"] = encoder_config["d_latent"]
elif encoder_config["d_model"] == 'kv':
encoder_config["d_model"] = model_cfg["d_feat"]
else:
raise ValueError(f"Unknown d_model: {encoder_config['d_model']}")
# # required for PerceiverTF with attention_to_channel option
# if self.encoder_type == "perceiver-tf":
# if encoder_config["attention_to_channel"] is True:
# encoder_config["kv_dim"] = model_cfg["d_feat"] # TODO: better readablity
# else:
# encoder_config["kv_dim"] = model_cfg["conv_out_channels"]
if isinstance(model_cfg["vocab_size"], str) and model_cfg["vocab_size"] == 'auto':
model_cfg["vocab_size"] = task_manager.num_tokens
if isinstance(model_cfg["num_max_positions"], str) and model_cfg["num_max_positions"] == 'auto':
model_cfg["num_max_positions"] = int(
max(model_cfg["feat_length"], model_cfg["event_length"]) + self.task_manager.max_task_token_length + 10)
# Pre-decoder
self.pre_decoder = nn.Sequential()
if self.encoder_type == "perceiver-tf" and self.decoder_type == "t5":
t, f, c = pre_enc_output_shape # perceiver-tf: (110, 128, 128) for 2s
encoder_output_shape = (t, encoder_config["num_latents"], encoder_config["d_latent"]) # (T, K, D_source)
decoder_input_shape = (t, decoder_config["d_model"]) # (T, D_target)
proj_layer = get_projection_layer(input_shape=encoder_output_shape,
output_shape=decoder_input_shape,
proj_type=self.pre_decoder_type)
self.pre_encoder_output_shape = pre_enc_output_shape
self.encoder_output_shape = encoder_output_shape
self.decoder_input_shape = decoder_input_shape
self.pre_decoder.append(proj_layer)
elif self.encoder_type in ["t5", "conformer"] and self.decoder_type == "t5":
pass
elif self.encoder_type == "perceiver-tf" and self.decoder_type == "multi-t5":
# NOTE: this is experiemental, only for multi-channel decoding with 13 classes
assert encoder_config["num_latents"] % decoder_config["num_channels"] == 0
encoder_output_shape = (encoder_config["num_latents"], encoder_config["d_model"])
decoder_input_shape = (decoder_config["num_channels"], decoder_config["d_model"])
proj_layer = get_multi_channel_projection_layer(input_shape=encoder_output_shape,
output_shape=decoder_input_shape,
proj_type=self.pre_decoder_type)
self.pre_decoder.append(proj_layer)
else:
raise NotImplementedError(
f"Encoder type {self.encoder_type} and decoder type {self.decoder_type} is not implemented yet.")
# Positional Encoding, Vocab, etc.
if self.encoder_type in ["t5", "conformer"]:
encoder_config["num_max_positions"] = decoder_config["num_max_positions"] = model_cfg["num_max_positions"]
else: # perceiver-tf uses separate positional encoding
encoder_config["num_max_positions"] = model_cfg["feat_length"]
decoder_config["num_max_positions"] = model_cfg["num_max_positions"]
encoder_config["vocab_size"] = decoder_config["vocab_size"] = model_cfg["vocab_size"]
# Print and save updated configs
self.audio_cfg = audio_cfg
self.model_cfg = model_cfg
self.shared_cfg = shared_cfg
self.save_hyperparameters()
if self.global_rank == 0:
print(self.hparams)
# Encoder and Decoder and LM-head
self.encoder = None
self.decoder = None
self.lm_head = LMHead(decoder_config, 1.0, model_cfg["tie_word_embeddings"])
self.embed_tokens = nn.Embedding(decoder_config["vocab_size"], decoder_config["d_model"])
self.embed_tokens.weight.data.normal_(mean=0.0, std=1.0)
self.shift_right_fn = None
self.set_encoder_decoder() # shift_right_fn is also set here
# Model as ModuleDict
# self.model = nn.ModuleDict({
# "pitchshift": self.pitchshift, # no grad; created in setup() only for training,
# and called by training_step()
# "spectrogram": self.spectrogram, # no grad
# "pre_encoder": self.pre_encoder,
# "encoder": self.encoder,
# "pre_decoder": self.pre_decoder,
# "decoder": self.decoder,
# "embed_tokens": self.embed_tokens,
# "lm_head": self.lm_head,
# })
# Tables (for logging)
columns = ['Ep', 'Track ID', 'Pred Events', 'Actual Events', 'Pred Notes', 'Actual Notes']
self.sample_table = wandb.Table(columns=columns)
# Output MIDI
if write_output_dir is not None:
if write_output_vocab is None:
from config.vocabulary import program_vocab_presets
self.midi_output_vocab = program_vocab_presets["gm_ext_plus"]
else:
self.midi_output_vocab = write_output_vocab
self.midi_output_inverse_vocab = create_inverse_vocab(self.midi_output_vocab)
def set_encoder_decoder(self) -> None:
"""Set encoder, decoder, lm_head and emb_tokens from self.model_cfg"""
# Generate and update T5Config
t5_basename = self.model_cfg["t5_basename"]
if t5_basename in T5_BASE_CFG.keys():
# Load from pre-defined config in config.py
t5_config = T5Config(**T5_BASE_CFG[t5_basename])
else:
# Load from HuggingFace hub
t5_config = T5Config.from_pretrained(t5_basename)
# Create encoder, decoder, lm_head and embed_tokens
if self.encoder_type == "t5":
self.encoder = T5EncoderYMT3(self.model_cfg["encoder"]["t5"], t5_config)
elif self.encoder_type == "perceiver-tf":
perceivertf_config = PerceiverTFConfig()
perceivertf_config.update(self.model_cfg["encoder"]["perceiver-tf"])
self.encoder = PerceiverTFEncoder(perceivertf_config)
elif self.encoder_type == "conformer":
conformer_config = ConformerYMT3Config()
conformer_config.update(self.model_cfg["encoder"]["conformer"])
self.encoder = ConformerYMT3Encoder(conformer_config)
if self.decoder_type == "t5":
self.decoder = T5DecoderYMT3(self.model_cfg["decoder"]["t5"], t5_config)
elif self.decoder_type == "multi-t5":
self.decoder = MultiChannelT5Decoder(self.model_cfg["decoder"]["multi-t5"], t5_config)
# `shift_right` function for decoding
self.shift_right_fn = self.decoder._shift_right
def setup(self, stage: str) -> None:
# Defining metrics
if self.hparams.eval_vocab is None:
extra_classes_per_dataset = [None]
else:
extra_classes_per_dataset = [
list(v.keys()) if v is not None else None for v in self.hparams.eval_vocab
] # e.g. [['Piano'], ['Guitar'], ['Piano'], ['Piano', 'Strings', 'Winds'], None]
# For direct addition of extra metrics using full metric name
extra_metrics = None
if self.hparams.add_melody_metric_to_singing is True:
extra_metrics = ["melody_rpa_Singing Voice", "melody_rca_Singing Voice", "melody_oa_Singing Voice"]
# Add pitch class metric
if self.hparams.add_pitch_class_metric is not None:
for sublist in extra_classes_per_dataset:
for name in self.hparams.add_pitch_class_metric:
if sublist is not None and name in sublist:
sublist += [name + "_pc"]
extra_classes_unique = list(
set(item for sublist in extra_classes_per_dataset if sublist is not None
for item in sublist)) # e.g. ['Strings', 'Winds', 'Guitar', 'Piano']
dm = self.trainer.datamodule
# Train/Vaidation-only
if stage == "fit":
self.val_metrics_macro = AMTMetrics(prefix=f'validation/macro_', extra_classes=extra_classes_unique)
self.val_metrics = nn.ModuleList() # val_metric is a list of AMTMetrics objects
for i in range(dm.num_val_dataloaders):
self.val_metrics.append(
AMTMetrics(prefix=f'validation/({dm.get_val_dataset_name(i)})',
extra_classes=extra_classes_per_dataset[i],
error_types=DECODING_ERR_TYPES))
# Add pitchshift layer
if self.shared_cfg["AUGMENTATION"]["train_pitch_shift_range"] in [None, [0, 0]]:
self.pitchshift = None
else:
# torchaudio pitchshifter requires a dummy input for initialization in DDP
input_shape = (self.shared_cfg["BSZ"]["train_local"], 1, self.audio_cfg["input_frames"])
self.pitchshift = PitchShiftLayer(
pshift_range=self.shared_cfg["AUGMENTATION"]["train_pitch_shift_range"],
expected_input_shape=input_shape,
device=self.device)
# Test-only
elif stage == "test":
# self.test_metrics_macro = AMTMetrics(
# prefix=f'test/macro_', extra_classes=extra_classes_unique)
self.test_metrics = nn.ModuleList()
for i in range(dm.num_test_dataloaders):
self.test_metrics.append(
AMTMetrics(prefix=f'test/({dm.get_test_dataset_name(i)})',
extra_classes=extra_classes_per_dataset[i],
extra_metrics=extra_metrics,
error_types=DECODING_ERR_TYPES))
# Test pitch shift layer: debug only
if self.test_pitch_shift_layer is not None:
self.test_pitch_shift_semitone = int(self.test_pitch_shift_layer)
self.pitchshift = PitchShiftLayer(
pshift_range=[self.test_pitch_shift_semitone, self.test_pitch_shift_semitone])
def configure_optimizers(self) -> None:
"""Configure optimizer and scheduler"""
optimizer, base_lr = get_optimizer(models_dict=self.named_parameters(),
optimizer_name=self.hparams.optimizer_name,
base_lr=self.hparams.base_lr,
weight_decay=self.hparams.weight_decay)
if self.hparams.optimizer_name.lower() == 'adafactor' and self.hparams.base_lr == None:
print("Using AdaFactor with auto learning rate and no scheduler")
return [optimizer]
if self.hparams.optimizer_name.lower() == 'dadaptadam':
print("Using dAdaptAdam with auto learning rate and no scheduler")
return [optimizer]
elif self.hparams.base_lr == None:
print(f"Using default learning rate {base_lr} of {self.hparams.optimizer_name} as base learning rate.")
self.hparams.base_lr = base_lr
scheduler_cfg = self.shared_cfg["LR_SCHEDULE"]
if self.hparams.max_steps != -1:
# overwrite total_steps
scheduler_cfg["total_steps"] = self.hparams.max_steps
_lr_scheduler = get_lr_scheduler(optimizer,
scheduler_name=self.hparams.scheduler_name,
base_lr=base_lr,
scheduler_cfg=scheduler_cfg)
lr_scheduler = {'scheduler': _lr_scheduler, 'interval': 'step', 'frequency': 1}
return [optimizer], [lr_scheduler]
def forward(
self,
x: torch.FloatTensor,
target_tokens: torch.LongTensor,
# task_tokens: Optional[torch.LongTensor] = None,
**kwargs) -> Dict:
""" Forward pass with teacher-forcing for training and validation.
Args:
x: (B, 1, T) waveform with default T=32767
target_tokens: (B, C, N) tokenized sequence of length N=event_length
task_tokens: (B, C, task_len) tokenized task
Returns:
{
'logits': (B, N + task_len + 1, vocab_size)
'loss': (1, )
}
NOTE: all the commented shapes are in the case of original MT3 setup.
"""
x = self.spectrogram(x) # mel-/spectrogram: (b, 256, 512) or (B, T, F)
x = self.pre_encoder(x) # projection to d_model: (B, 256, 512)
# TODO: task_cond_encoder would not work properly because of 3-d task_tokens
# if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_encoder is True:
# # append task embedding to encoder input
# task_embed = self.embed_tokens(task_tokens) # (B, task_len, 512)
# x = torch.cat([task_embed, x], dim=1) # (B, task_len + 256, 512)
enc_hs = self.encoder(inputs_embeds=x)["last_hidden_state"] # (B, T', D)
enc_hs = self.pre_decoder(enc_hs) # (B, T', D) or (B, K, T, D)
# if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_decoder is True:
# # append task token to decoder input and output label
# labels = torch.cat([task_tokens, target_tokens], dim=2) # (B, C, task_len + N)
# else:
# labels = target_tokens # (B, C, N)
labels = target_tokens # (B, C, N)
if labels.shape[1] == 1: # for single-channel decoders, e.g. t5.
labels = labels.squeeze(1) # (B, N)
dec_input_ids = self.shift_right_fn(labels) # t5:(B, N), multi-t5:(B, C, N)
dec_inputs_embeds = self.embed_tokens(dec_input_ids) # t5:(B, N, D), multi-t5:(B, C, N, D)
dec_hs, _ = self.decoder(inputs_embeds=dec_inputs_embeds, encoder_hidden_states=enc_hs, return_dict=False)
if self.model_cfg["tie_word_embeddings"] is True:
dec_hs = dec_hs * (self.model_cfg["decoder"][self.decoder_type]["d_model"]**-0.5)
logits = self.lm_head(dec_hs)
loss = None
labels = labels.masked_fill(labels == 0, value=-100) # ignore pad tokens for loss
loss_fct = CrossEntropyLoss(ignore_index=-100)
loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
return {"logits": logits, "loss": loss}
def inference(self,
x: torch.FloatTensor,
task_tokens: Optional[torch.LongTensor] = None,
max_token_length: Optional[int] = None,
**kwargs: Any) -> torch.Tensor:
""" Inference from audio batch by cached autoregressive decoding.
Args:
x: (b, 1, t) waveform with t=32767
task_token: (b, c, task_len) tokenized task. If None, will not append task embeddings (from task_tokens) to input.
max_length: Maximum length of generated sequence. If None, self.max_total_token_length.
**kwargs: https://huggingface.co/docs/transformers/v4.27.2/en/main_classes/text_generation#transformers.GenerationMixin.generate
Returns:
res_tokens: (b, n) resulting tokenized sequence of variable length < max_length
"""
if self.test_pitch_shift_layer is not None:
x_ps = self.pitchshift(x, self.test_pitch_shift_semitone)
x = x_ps
# From spectrogram to pre-decoder is the same pipeline as in forward()
x = self.spectrogram(x) # mel-/spectrogram: (b, 256, 512) or (B, T, F)
x = self.pre_encoder(x) # projection to d_model: (B, 256, 512)
if task_tokens is not None and task_tokens.numel() > 0 and self.use_task_cond_encoder is True:
# append task embedding to encoder input
task_embed = self.embed_tokens(task_tokens) # (B, task_len, 512)
x = torch.cat([task_embed, x], dim=1) # (B, task_len + 256, 512)
enc_hs = self.encoder(inputs_embeds=x)["last_hidden_state"] # (B, task_len + 256, 512)
enc_hs = self.pre_decoder(enc_hs) # (B, task_len + 256, 512)
# Cached-autoregressive decoding with task token (can be None) as prefix
if max_token_length is None:
max_token_length = self.max_total_token_length
pred_ids = task_cond_dec_generate(decoder=self.decoder,
decoder_type=self.decoder_type,
embed_tokens=self.embed_tokens,
lm_head=self.lm_head,
encoder_hidden_states=enc_hs,
shift_right_fn=self.shift_right_fn,
prefix_ids=task_tokens,
max_length=max_token_length) # (B, task_len + N) or (B, C, task_len + N)
if pred_ids.dim() == 2:
pred_ids = pred_ids.unsqueeze(1) # (B, 1, task_len + N)
if self.test_pitch_shift_layer is None:
return pred_ids
else:
return pred_ids, x_ps
def inference_file(
self,
bsz: int,
audio_segments: torch.FloatTensor, # (n_items, 1, segment_len): from a single file
note_token_array: Optional[torch.LongTensor] = None,
task_token_array: Optional[torch.LongTensor] = None,
# subtask_key: Optional[str] = "default"
) -> Tuple[List[np.ndarray], Optional[torch.Tensor]]:
""" Inference from audio batch by autoregressive decoding:
Args:
bsz: batch size
audio_segments: (n_items, 1, segment_len): segmented audio from a single file
note_token_array: (n_items, max_token_len): Optional. If token_array is None, will not return loss.
subtask_key: (str): If None, not using subtask prefix. By default, using "default" defined in task manager.
"""
# if subtask_key is not None:
# _subtask_token = torch.LongTensor(
# self.task_manager.get_eval_subtask_prefix_dict()[subtask_key]).to(self.device)
n_items = audio_segments.shape[0]
loss = 0.
pred_token_array_file = [] # each element is (B, C, L) np.ndarray
x_ps_concat = []
for i in range(0, n_items, bsz):
if i + bsz > n_items: # last batch can be smaller
x = audio_segments[i:n_items].to(self.device)
# if subtask_key is not None:
# b = n_items - i # bsz for the last batch
# task_tokens = _subtask_token.expand((b, -1)) # (b, task_len)
if note_token_array is not None:
target_tokens = note_token_array[i:n_items].to(self.device)
if task_token_array is not None and task_token_array.numel() > 0:
task_tokens = task_token_array[i:n_items].to(self.device)
else:
task_tokens = None
else:
x = audio_segments[i:i + bsz].to(self.device) # (bsz, 1, segment_len)
# if subtask_key is not None:
# task_tokens = _subtask_token.expand((bsz, -1)) # (bsz, task_len)
if note_token_array is not None:
target_tokens = note_token_array[i:i + bsz].to(self.device) # (bsz, token_len)
if task_token_array is not None and task_token_array.numel() > 0:
task_tokens = task_token_array[i:i + bsz].to(self.device)
else:
task_tokens = None
# token prediction (fast-autoregressive decoding)
# if subtask_key is not None:
# preds = self.inference(x, task_tokens).detach().cpu().numpy()
# else:
# preds = self.inference(x).detach().cpu().numpy()
if self.test_pitch_shift_layer is not None: # debug only
preds, x_ps = self.inference(x, task_tokens)
preds = preds.detach().cpu().numpy()
x_ps_concat.append(x_ps.detach().cpu())
else:
preds = self.inference(x, task_tokens).detach().cpu().numpy()
if len(preds) != len(x):
raise ValueError(f'preds: {len(preds)}, x: {len(x)}')
pred_token_array_file.append(preds)
# validation loss (by teacher forcing)
if note_token_array is not None:
loss_weight = x.shape[0] / n_items
loss += self(x, target_tokens)['loss'] * loss_weight
# loss += self(x, target_tokens, task_tokens)['loss'] * loss_weight
else:
loss = None
if self.test_pitch_shift_layer is not None: # debug only
if self.hparams.write_output_dir is not None:
x_ps_concat = torch.cat(x_ps_concat, dim=0)
return pred_token_array_file, loss, x_ps_concat.flatten().unsqueeze(0)
else:
return pred_token_array_file, loss
def training_step(self, batch, batch_idx) -> torch.Tensor:
# batch: {
# 'dataset1': [Tuple[audio_segments(b, 1, t), tokens(b, max_token_len), ...]]
# 'dataset2': [Tuple[audio_segments(b, 1, t), tokens(b, max_token_len), ...]]
# 'dataset3': ...
# }
audio_segments, note_tokens, pshift_steps = [torch.cat(t, dim=0) for t in zip(*batch.values())]
if self.pitchshift is not None:
# Pitch shift
n_groups = len(batch)
audio_segments = torch.chunk(audio_segments, n_groups, dim=0)
pshift_steps = torch.chunk(pshift_steps, n_groups, dim=0)
for p in pshift_steps:
assert p.eq(p[0]).all().item()
audio_segments = torch.cat([self.pitchshift(a, p[0].item()) for a, p in zip(audio_segments, pshift_steps)],
dim=0)
loss = self(audio_segments, note_tokens)['loss']
self.log('train_loss',
loss,
on_step=True,
on_epoch=True,
prog_bar=True,
batch_size=note_tokens.shape[0],
sync_dist=True)
# print('lr', self.trainer.optimizers[0].param_groups[0]['lr'])
return loss
def validation_step(self, batch, batch_idx, dataloader_idx=0) -> Dict:
# File-wise validation
if self.task_manager.num_decoding_channels == 1:
bsz = self.shared_cfg["BSZ"]["validation"]
else:
bsz = self.shared_cfg["BSZ"]["validation"] // self.task_manager.num_decoding_channels * 3
# audio_segments, notes_dict, note_token_array, task_token_array = batch
audio_segments, notes_dict, note_token_array = batch
task_token_array = None
# Loop through the tensor in chunks of bsz (=subbsz actually)
n_items = audio_segments.shape[0]
start_secs_file = [32767 * i / 16000 for i in range(n_items)]
with Timer() as t:
pred_token_array_file, loss = self.inference_file(bsz, audio_segments, note_token_array, task_token_array)
"""
notes_dict: # Ground truth notes
{
'mtrack_id': int,
'program': List[int],
'is_drum': bool,
'duration_sec': float,
'notes': List[Note],
}
"""
# Process a list of channel-wise token arrays for a file
num_channels = self.task_manager.num_decoding_channels
pred_notes_in_file = []
n_err_cnt = Counter()
for ch in range(num_channels):
pred_token_array_ch = [arr[:, ch, :] for arr in pred_token_array_file] # (B, L)
zipped_note_events_and_tie, list_events, ne_err_cnt = self.task_manager.detokenize_list_batches(
pred_token_array_ch, start_secs_file, return_events=True)
pred_notes_ch, n_err_cnt_ch = merge_zipped_note_events_and_ties_to_notes(zipped_note_events_and_tie)
pred_notes_in_file.append(pred_notes_ch)
n_err_cnt += n_err_cnt_ch
pred_notes = mix_notes(pred_notes_in_file) # This is the mixed notes from all channels
if self.hparams.write_output_dir is not None:
track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
# write_model_output_as_npy(zipped_note_events_and_tie, self.hparams.write_output_dir,
# track_info)
write_model_output_as_midi(pred_notes,
self.hparams.write_output_dir,
track_info,
self.midi_output_inverse_vocab,
output_dir_suffix=str(dataset_info) + '_' +
str(self.hparams.eval_subtask_key))
# generate sample text to display in log table
# pred_events_text = [str([list_events[0][:200]])]
# pred_notes_text = [str([pred_notes[:200]])]
# this is local GPU metric per file, not global metric in DDP
drum_metric, non_drum_metric, instr_metric = compute_track_metrics(
pred_notes,
notes_dict['notes'],
eval_vocab=self.hparams.eval_vocab[dataloader_idx],
eval_drum_vocab=self.hparams.eval_drum_vocab,
onset_tolerance=self.hparams.onset_tolerance,
add_pitch_class_metric=self.hparams.add_pitch_class_metric)
self.val_metrics[dataloader_idx].bulk_update(drum_metric)
self.val_metrics[dataloader_idx].bulk_update(non_drum_metric)
self.val_metrics[dataloader_idx].bulk_update(instr_metric)
self.val_metrics_macro.bulk_update(drum_metric)
self.val_metrics_macro.bulk_update(non_drum_metric)
self.val_metrics_macro.bulk_update(instr_metric)
# Log sample table: predicted notes and ground truth notes
# if batch_idx in (0, 1) and self.global_rank == 0:
# actual_notes_text = [str([notes_dict['notes'][:200]])]
# actual_tokens = token_array[0, :200].detach().cpu().numpy().tolist()
# actual_events_text = [str(self.tokenizer._decode(actual_tokens))]
# track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")]
# self.sample_table.add_data(self.current_epoch, track_info, pred_events_text,
# actual_events_text, pred_notes_text, actual_notes_text)
# self.logger.log_table('Samples', self.sample_table.columns, self.sample_table.data)
decoding_time_sec = t.elapsed_time()
self.log('val_loss', loss, prog_bar=True, batch_size=n_items, sync_dist=True)
# self.val_metrics[dataloader_idx].bulk_update_errors({'decoding_time': decoding_time_sec})
def on_validation_epoch_end(self) -> None:
for val_metrics in self.val_metrics:
self.log_dict(val_metrics.bulk_compute(), sync_dist=True)
val_metrics.bulk_reset()
self.log_dict(self.val_metrics_macro.bulk_compute(), sync_dist=True)
self.val_metrics_macro.bulk_reset()
def test_step(self, batch, batch_idx, dataloader_idx=0) -> Dict:
# File-wise evaluation
if self.task_manager.num_decoding_channels == 1:
bsz = self.shared_cfg["BSZ"]["validation"]
else:
bsz = self.shared_cfg["BSZ"]["validation"] // self.task_manager.num_decoding_channels * 3
# audio_segments, notes_dict, note_token_array, task_token_array = batch
audio_segments, notes_dict, note_token_array = batch
task_token_array = None
# Test pitch shift layer: debug only
if self.test_pitch_shift_layer is not None and self.test_pitch_shift_semitone != 0:
for n in notes_dict['notes']:
if n.is_drum == False:
n.pitch = n.pitch + self.test_pitch_shift_semitone
# Loop through the tensor in chunks of bsz (=subbsz actually)
n_items = audio_segments.shape[0]
start_secs_file = [32767 * i / 16000 for i in range(n_items)]
if self.test_pitch_shift_layer is not None and self.hparams.write_output_dir is not None:
pred_token_array_file, loss, x_ps = self.inference_file(bsz, audio_segments, None, None)
else:
pred_token_array_file, loss = self.inference_file(bsz, audio_segments, None, None)
if len(pred_token_array_file) > 0:
# Process a list of channel-wise token arrays for a file
num_channels = self.task_manager.num_decoding_channels
pred_notes_in_file = []
n_err_cnt = Counter()
for ch in range(num_channels):
pred_token_array_ch = [arr[:, ch, :] for arr in pred_token_array_file] # (B, L)
zipped_note_events_and_tie, list_events, ne_err_cnt = self.task_manager.detokenize_list_batches(
pred_token_array_ch, start_secs_file, return_events=True)
pred_notes_ch, n_err_cnt_ch = merge_zipped_note_events_and_ties_to_notes(zipped_note_events_and_tie)
pred_notes_in_file.append(pred_notes_ch)
n_err_cnt += n_err_cnt_ch
pred_notes = mix_notes(pred_notes_in_file) # This is the mixed notes from all channels
if self.test_pitch_shift_layer is not None and self.hparams.write_output_dir is not None:
# debug only
wav_output_dir = os.path.join(self.hparams.write_output_dir, f"model_output_{dataset_info}")
os.makedirs(wav_output_dir, exist_ok=True)
wav_output_file = os.path.join(wav_output_dir, f"{track_info}_ps_{self.test_pitch_shift_semitone}.wav")
torchaudio.save(wav_output_file, x_ps.squeeze(1), 16000, bits_per_sample=16)
drum_metric, non_drum_metric, instr_metric = compute_track_metrics(
pred_notes,
notes_dict['notes'],
eval_vocab=self.hparams.eval_vocab[dataloader_idx],
eval_drum_vocab=self.hparams.eval_drum_vocab,
onset_tolerance=self.hparams.onset_tolerance,
add_pitch_class_metric=self.hparams.add_pitch_class_metric,
add_melody_metric=['Singing Voice'] if self.hparams.add_melody_metric_to_singing else None,
add_frame_metric=True,
add_micro_metric=True,
add_multi_f_metric=True)
if self.hparams.write_output_dir is not None and self.global_rank == 0:
# write model output to file
track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
f_score = f"OnF{non_drum_metric['onset_f']:.2f}_MulF{instr_metric['multi_f']:.2f}"
write_model_output_as_midi(pred_notes,
self.hparams.write_output_dir,
track_info,
self.midi_output_inverse_vocab,
output_dir_suffix=str(dataset_info) + '_' +
str(self.hparams.eval_subtask_key) + '_' + f_score)
write_err_cnt_as_json(track_info, self.hparams.write_output_dir,
str(dataset_info) + '_' + str(self.hparams.eval_subtask_key) + '_' + f_score,
n_err_cnt, ne_err_cnt)
# Test with optimal octave shift
if self.hparams.test_optimal_octave_shift:
track_info = [notes_dict[k] for k in notes_dict.keys() if k.endswith("_id")][0]
dataset_info = [k for k in notes_dict.keys() if k.endswith('_id')][0][:-3]
score = [instr_metric['onset_f_Bass']]
ref_notes_plus = []
ref_notes_minus = []
for note in notes_dict['notes']:
if note.is_drum == True:
ref_notes_plus.append(note)
ref_notes_minus.append(note)
else:
ref_notes_plus.append(
Note(is_drum=note.is_drum,
program=note.program,
onset=note.onset,
offset=note.offset,
pitch=note.pitch + 12,
velocity=note.velocity))
ref_notes_minus.append(
Note(is_drum=note.is_drum,
program=note.program,
onset=note.onset,
offset=note.offset,
pitch=note.pitch - 12,
velocity=note.velocity))
drum_metric_plus, non_drum_metric_plus, instr_metric_plus = compute_track_metrics(
pred_notes,
ref_notes_plus,
eval_vocab=self.hparams.eval_vocab[dataloader_idx],
eval_drum_vocab=self.hparams.eval_drum_vocab,
onset_tolerance=self.hparams.onset_tolerance,
add_pitch_class_metric=self.hparams.add_pitch_class_metric)
drum_metric_minus, non_drum_metric_minus, instr_metric_minus = compute_track_metrics(
ref_notes_minus,
notes_dict['notes'],
eval_vocab=self.hparams.eval_vocab[dataloader_idx],
eval_drum_vocab=self.hparams.eval_drum_vocab,
onset_tolerance=self.hparams.onset_tolerance,
add_pitch_class_metric=self.hparams.add_pitch_class_metric)
score.append(instr_metric_plus['onset_f_Bass'])
score.append(instr_metric_minus['onset_f_Bass'])
max_index = score.index(max(score))
if max_index == 0:
print(f"ZERO: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
elif max_index == 1:
# plus
instr_metric['onset_f_Bass'] = instr_metric_plus['onset_f_Bass']
print(f"PLUS: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
write_model_output_as_midi(ref_notes_plus,
self.hparams.write_output_dir,
track_info + '_ref_octave_plus',
self.midi_output_inverse_vocab,
output_dir_suffix=str(dataset_info) + '_' +
str(self.hparams.eval_subtask_key))
else:
# minus
instr_metric['onset_f_Bass'] = instr_metric_minus['onset_f_Bass']
print(f"MINUS: {track_info}, z/p/m: {score[0]:.2f}/{score[1]:.2f}/{score[2]:.2f}")
write_model_output_as_midi(ref_notes_minus,
self.hparams.write_output_dir,
track_info + '_ref_octave_minus',
self.midi_output_,
output_dir_suffix=str(dataset_info) + '_' +
str(self.hparams.eval_subtask_key))
self.test_metrics[dataloader_idx].bulk_update(drum_metric)
self.test_metrics[dataloader_idx].bulk_update(non_drum_metric)
self.test_metrics[dataloader_idx].bulk_update(instr_metric)
# self.test_metrics_macro.bulk_update(drum_metric)
# self.test_metrics_macro.bulk_update(non_drum_metric)
# self.test_metrics_macro.bulk_update(instr_metric)
def on_test_epoch_end(self) -> None:
# all_gather is done seeminglesly by torchmetrics
for test_metrics in self.test_metrics:
self.log_dict(test_metrics.bulk_compute(), sync_dist=True)
test_metrics.bulk_reset()
# self.log_dict(self.test_metrics_macro.bulk_compute(), sync_dist=True)
# self.test_metrics_macro.bulk_reset()
def test_case_forward_mt3():
import torch
from config.config import audio_cfg, model_cfg, shared_cfg
from model.ymt3 import YourMT3
model = YourMT3()
model.eval()
x = torch.randn(2, 1, 32767)
labels = torch.randint(0, 596, (2, 1, 1024), requires_grad=False) # (B, C=1, T)
task_tokens = torch.LongTensor([])
output = model.forward(x, labels, task_tokens)
logits, loss = output['logits'], output['loss']
assert logits.shape == (2, 1024, 596) # (B, N, vocab_size)
def test_case_inference_mt3():
import torch
from config.config import audio_cfg, model_cfg, shared_cfg
from model.ymt3 import YourMT3
model_cfg["num_max_positions"] = 1024 + 3 + 1
model = YourMT3(model_cfg=model_cfg)
model.eval()
x = torch.randn(2, 1, 32767)
task_tokens = torch.randint(0, 596, (2, 3), requires_grad=False)
pred_ids = model.inference(x, task_tokens, max_token_length=10) # (2, 3, 9) (B, C, L-task_len)
# TODO: need to check the length of pred_ids when task_tokens is not None
def test_case_forward_enc_perceiver_tf_dec_t5():
import torch
from model.ymt3 import YourMT3
from config.config import audio_cfg, model_cfg, shared_cfg
model_cfg["encoder_type"] = "perceiver-tf"
audio_cfg["codec"] = "spec"
audio_cfg["hop_length"] = 300
model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
model.eval()
x = torch.randn(2, 1, 32767)
labels = torch.randint(0, 596, (2, 1, 1024), requires_grad=False)
# forward
output = model.forward(x, labels)
logits, loss = output['logits'], output['loss'] # logits: (2, 1024, 596) (B, N, vocab_size)
# inference
pred_ids = model.inference(x, None, max_token_length=3) # (2, 1, 3) (B, C, L)
def test_case_forward_enc_conformer_dec_t5():
import torch
from model.ymt3 import YourMT3
from config.config import audio_cfg, model_cfg, shared_cfg
model_cfg["encoder_type"] = "conformer"
audio_cfg["codec"] = "melspec"
audio_cfg["hop_length"] = 128
model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
model.eval()
x = torch.randn(2, 1, 32767)
labels = torch.randint(0, 596, (2, 1024), requires_grad=False)
# forward
output = model.forward(x, labels)
logits, loss = output['logits'], output['loss'] # logits: (2, 1024, 596) (B, N, vocab_size)
# inference
pred_ids = model.inference(x, None, 20) # (2, 1, 20) (B, C, L)
def test_case_enc_perceiver_tf_dec_multi_t5():
import torch
from model.ymt3 import YourMT3
from config.config import audio_cfg, model_cfg, shared_cfg
model_cfg["encoder_type"] = "perceiver-tf"
model_cfg["decoder_type"] = "multi-t5"
model_cfg["encoder"]["perceiver-tf"]["attention_to_channel"] = True
model_cfg["encoder"]["perceiver-tf"]["num_latents"] = 26
audio_cfg["codec"] = "spec"
audio_cfg["hop_length"] = 300
model = YourMT3(audio_cfg=audio_cfg, model_cfg=model_cfg)
model.eval()
x = torch.randn(2, 1, 32767)
labels = torch.randint(0, 596, (2, 13, 200), requires_grad=False) # (B, C, T)
# x = model.spectrogram(x)
# x = model.pre_encoder(x) # (2, 110, 128, 128) (B, T, C, D)
# enc_hs = model.encoder(inputs_embeds=x)["last_hidden_state"] # (2, 110, 128, 128) (B, T, C, D)
# enc_hs = model.pre_decoder(enc_hs) # (2, 13, 110, 512) (B, C, T, D)
# dec_input_ids = model.shift_right_fn(labels) # (2, 13, 200) (B, C, T)
# dec_inputs_embeds = model.embed_tokens(dec_input_ids) # (2, 13, 200, 512) (B, C, T, D)
# dec_hs, _ = model.decoder(
# inputs_embeds=dec_inputs_embeds, encoder_hidden_states=enc_hs, return_dict=False)
# logits = model.lm_head(dec_hs) # (2, 13, 200, 596) (B, C, T, vocab_size)
# forward
x = torch.randn(2, 1, 32767)
labels = torch.randint(0, 596, (2, 13, 200), requires_grad=False) # (B, C, T)
output = model.forward(x, labels)
logits, loss = output['logits'], output['loss'] # (2, 13, 200, 596) (B, C, T, vocab_size)
# inference
model.max_total_token_length = 123 # to save time..
pred_ids = model.inference(x, None) # (2, 13, 123) (B, C, L)