train.py

import argparse
import os
import random
from pathlib import Path
from typing import Union

import lightning as pl
import numpy as np
import torch
import torch.nn.functional as F
from lightning import Trainer
from lightning.fabric.utilities import rank_zero_only
from lightning.pytorch.callbacks import ModelCheckpoint
from peft import LoraConfig, TaskType
from safetensors.torch import save_file as safe_save_file
from torch import optim
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import Dataset, DataLoader

import MIDI
from midi_model import MIDIModel, MIDIModelConfig, config_name_list
from midi_tokenizer import MIDITokenizerV1, MIDITokenizerV2

EXTENSION = [".mid", ".midi"]


def file_ext(fname):
    return os.path.splitext(fname)[1].lower()


class MidiDataset(Dataset):
    def __init__(self, midi_list, tokenizer: Union[MIDITokenizerV1, MIDITokenizerV2], max_len=2048, min_file_size=3000,
                 max_file_size=384000,
                 aug=True, check_quality=False, rand_start=True):

        self.tokenizer = tokenizer
        self.midi_list = midi_list
        self.max_len = max_len
        self.min_file_size = min_file_size
        self.max_file_size = max_file_size
        self.aug = aug
        self.check_quality = check_quality
        self.rand_start = rand_start

    def __len__(self):
        return len(self.midi_list)

    def load_midi(self, index):
        path = self.midi_list[index]
        try:
            with open(path, 'rb') as f:
                datas = f.read()
            if len(datas) > self.max_file_size:  # large midi file will spend too much time to load
                raise ValueError("file too large")
            elif len(datas) < self.min_file_size:
                raise ValueError("file too small")
            mid = MIDI.midi2score(datas)
            if max([0] + [len(track) for track in mid[1:]]) == 0:
                raise ValueError("empty track")
            mid = self.tokenizer.tokenize(mid)
            if self.check_quality and not self.tokenizer.check_quality(mid)[0]:
                raise ValueError("bad quality")
            if self.aug:
                mid = self.tokenizer.augment(mid)
        except Exception:
            mid = self.load_midi(random.randint(0, self.__len__() - 1))
        return mid

    def __getitem__(self, index):
        mid = self.load_midi(index)
        mid = np.asarray(mid, dtype=np.int16)
        # if mid.shape[0] < self.max_len:
        #     mid = np.pad(mid, ((0, self.max_len - mid.shape[0]), (0, 0)),
        #                  mode="constant", constant_values=self.tokenizer.pad_id)
        if self.rand_start:
            start_idx = random.randrange(0, max(1, mid.shape[0] - self.max_len))
            start_idx = random.choice([0, start_idx])
        else:
            max_start = max(1, mid.shape[0] - self.max_len)
            start_idx = (index * (max_start // 8)) % max_start
        mid = mid[start_idx: start_idx + self.max_len]
        mid = mid.astype(np.int64)
        mid = torch.from_numpy(mid)
        return mid

    def collate_fn(self, batch):
        max_len = max([len(mid) for mid in batch])
        batch = [F.pad(mid, (0, 0, 0, max_len - mid.shape[0]), mode="constant", value=self.tokenizer.pad_id) for mid in batch]
        batch = torch.stack(batch)
        return batch


def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
    """ Create a schedule with a learning rate that decreases linearly after
    linearly increasing during a warmup period.
    """

    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))

    return LambdaLR(optimizer, lr_lambda, last_epoch)


class TrainMIDIModel(MIDIModel, pl.LightningModule):
    def __init__(self, config: MIDIModelConfig,
                 lr=2e-4, weight_decay=0.01, warmup=1e3, max_step=1e6, sample_seq=False,
                 gen_example_interval=1, example_batch=8):
        super(TrainMIDIModel, self).__init__(config)
        self.lr = lr
        self.weight_decay = weight_decay
        self.warmup = warmup
        self.max_step = max_step
        self.sample_seq = sample_seq
        self.gen_example_interval = gen_example_interval
        self.example_batch = example_batch
        self.last_save_step = 0
        self.gen_example_count = 0

    def configure_optimizers(self):
        param_optimizer = list(self.named_parameters())
        no_decay = ['bias', 'norm']  # no decay for bias and Norm
        optimizer_grouped_parameters = [
            {
                'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
                'weight_decay': self.weight_decay},
            {
                'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
                'weight_decay': 0.0
            }
        ]
        optimizer = optim.AdamW(
            optimizer_grouped_parameters,
            lr=self.lr,
            betas=(0.9, 0.99),
            eps=1e-08,
        )
        lr_scheduler = get_linear_schedule_with_warmup(
            optimizer=optimizer,
            num_warmup_steps=self.warmup,
            num_training_steps=self.max_step,
        )
        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": lr_scheduler,
                "interval": "step",
                "frequency": 1
            }
        }

    def compute_accuracy(self, logits, labels):
        out = torch.argmax(logits, dim=-1)
        out = out.flatten()
        labels = labels.flatten()

        mask = (labels != self.tokenizer.pad_id)
        out = out[mask]
        labels = labels[mask]

        num_right = (out == labels)
        num_right = torch.sum(num_right).type(torch.float32)
        acc = num_right / len(labels)

        return acc

    def training_step(self, batch, batch_idx):
        x = batch[:, :-1].contiguous()  # (batch_size, midi_sequence_length, token_sequence_length)
        y = batch[:, 1:].contiguous()
        hidden = self.forward(x)
        if self.sample_seq:  # to reduce vram
            rand_idx = [-1] + random.sample(list(range(y.shape[1] - 2)), min(127, (y.shape[1] - 2) // 2))
            hidden = hidden[:, rand_idx]
            y = y[:, rand_idx]
        hidden = hidden.reshape(-1, hidden.shape[-1])
        y = y.reshape(-1, y.shape[-1])  # (batch_size*midi_sequence_length, token_sequence_length)
        x = y[:, :-1]
        logits = self.forward_token(hidden, x)
        loss = F.cross_entropy(
            logits.view(-1, self.tokenizer.vocab_size),
            y.view(-1),
            reduction="mean",
            ignore_index=self.tokenizer.pad_id
        )
        self.log("train/loss", loss)
        self.log("train/lr", self.lr_schedulers().get_last_lr()[0])
        return loss

    def validation_step(self, batch, batch_idx):
        x = batch[:, :-1].contiguous()  # (batch_size, midi_sequence_length, token_sequence_length)
        y = batch[:, 1:].contiguous()
        hidden = self.forward(x)
        hidden = hidden.reshape(-1, hidden.shape[-1])
        y = y.reshape(-1, y.shape[-1])  # (batch_size*midi_sequence_length, token_sequence_length)
        x = y[:, :-1]
        logits = self.forward_token(hidden, x)
        loss = F.cross_entropy(
            logits.view(-1, self.tokenizer.vocab_size),
            y.view(-1),
            reduction="mean",
            ignore_index=self.tokenizer.pad_id
        )
        acc = self.compute_accuracy(logits, y)
        self.log_dict({"val/loss": loss, "val/acc": acc}, sync_dist=True)
        return loss

    @rank_zero_only
    def gen_example(self, save_dir):
        base_dir = f"{save_dir}/sample/{self.global_step}"
        if not os.path.exists(base_dir):
            Path(base_dir).mkdir(parents=True)
        midis = self.generate(batch_size=self.example_batch)
        midis = [self.tokenizer.detokenize(midi) for midi in midis]
        imgs = [self.tokenizer.midi2img(midi) for midi in midis]
        for i, (img, midi) in enumerate(zip(imgs, midis)):
            img.save(f"{base_dir}/0_{i}.png")
            with open(f"{base_dir}/0_{i}.mid", 'wb') as f:
                f.write(MIDI.score2midi(midi))
        prompt = val_dataset.load_midi(random.randint(0, len(val_dataset) - 1))
        prompt = np.asarray(prompt, dtype=np.int16)
        ori = prompt[:512]
        img = self.tokenizer.midi2img(self.tokenizer.detokenize(ori))
        img.save(f"{base_dir}/1_ori.png")
        prompt = prompt[:256].astype(np.int64)
        midis = self.generate(prompt, batch_size=self.example_batch)
        midis = [self.tokenizer.detokenize(midi) for midi in midis]
        imgs = [self.tokenizer.midi2img(midi) for midi in midis]
        for i, (img, midi) in enumerate(zip(imgs, midis)):
            img.save(f"{base_dir}/1_{i}.png")
            with open(f"{base_dir}/1_{i}.mid", 'wb') as f:
                f.write(MIDI.score2midi(midi))

    @rank_zero_only
    def save_peft(self, save_dir):
        adapter_name = self.active_adapters()[0]
        adapter_config = self.peft_config[adapter_name]
        if not os.path.exists(save_dir):
            os.makedirs(save_dir, exist_ok=True)
        adapter_config.save_pretrained(save_dir)
        adapter_state_dict = self.get_adapter_state_dict(adapter_name)
        safe_save_file(adapter_state_dict,
                       os.path.join(save_dir, "adapter_model.safetensors"),
                       metadata={"format": "pt"})

    def on_save_checkpoint(self, checkpoint):
        if self.global_step == self.last_save_step:
            return
        self.last_save_step = self.global_step
        trainer = self.trainer
        if len(trainer.loggers) > 0:
            if trainer.loggers[0].save_dir is not None:
                save_dir = trainer.loggers[0].save_dir
            else:
                save_dir = trainer.default_root_dir
            name = trainer.loggers[0].name
            version = trainer.loggers[0].version
            version = version if isinstance(version, str) else f"version_{version}"
            save_dir = os.path.join(save_dir, str(name), version)
        else:
            save_dir = trainer.default_root_dir
        self.config.save_pretrained(os.path.join(save_dir, "checkpoints"))
        if self._hf_peft_config_loaded:
            self.save_peft(os.path.join(save_dir, "lora"))
        self.gen_example_count += 1
        if self.gen_example_interval>0 and self.gen_example_count % self.gen_example_interval == 0:
            try:
                self.gen_example(save_dir)
            except Exception as e:
                print(e)


def get_midi_list(path):
    all_files = {
        os.path.join(root, fname)
        for root, _dirs, files in os.walk(path)
        for fname in files
    }
    all_midis = sorted(
        fname for fname in all_files if file_ext(fname) in EXTENSION
    )
    return all_midis


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    # model args
    parser.add_argument(
        "--resume", type=str, default="", help="resume training from ckpt"
    )
    parser.add_argument(
        "--ckpt", type=str, default="", help="load ckpt"
    )
    parser.add_argument(
        "--config", type=str, default="tv2o-medium", help="model config name or file"
    )
    parser.add_argument(
        "--task", type=str, default="train", choices=["train", "lora"], help="Full train or lora"
    )

    # dataset args
    parser.add_argument(
        "--data", type=str, default="data", help="dataset path"
    )
    parser.add_argument(
        "--data-val-split",
        type=int,
        default=128,
        help="the number of midi files divided into the validation set",
    )
    parser.add_argument(
        "--max-len",
        type=int,
        default=2048,
        help="max seq length for training",
    )
    parser.add_argument(
        "--quality", action="store_true", default=False, help="check dataset quality"
    )

    # training args
    parser.add_argument("--seed", type=int, default=0, help="seed")
    parser.add_argument("--lr", type=float, default=1e-4, help="learning rate")
    parser.add_argument("--weight-decay", type=float, default=0.01, help="weight decay")
    parser.add_argument("--warmup-step", type=int, default=1e2, help="warmup step")
    parser.add_argument("--max-step", type=int, default=1e6, help="max training step")
    parser.add_argument("--grad-clip", type=float, default=1.0, help="gradient clip val")
    parser.add_argument(
        "--sample-seq", action="store_true", default=False, help="sample midi seq to reduce vram"
    )
    parser.add_argument(
        "--gen-example-interval", type=int, default=1, help="generate example interval. set 0 to disable"
    )
    parser.add_argument(
        "--batch-size-train", type=int, default=2, help="batch size for training"
    )
    parser.add_argument(
        "--batch-size-val", type=int, default=2, help="batch size for val"
    )
    parser.add_argument(
        "--batch-size-gen-example", type=int, default=8, help="batch size for generate example"
    )
    parser.add_argument(
        "--workers-train",
        type=int,
        default=4,
        help="workers num for training dataloader",
    )
    parser.add_argument(
        "--workers-val",
        type=int,
        default=4,
        help="workers num for validation dataloader",
    )
    parser.add_argument(
        "--acc-grad", type=int, default=2, help="gradient accumulation"
    )
    parser.add_argument(
        "--accelerator",
        type=str,
        default="gpu",
        choices=["cpu", "gpu", "tpu", "ipu", "hpu", "auto"],
        help="accelerator",
    )
    parser.add_argument(
        "--precision",
        type=str,
        default="bf16-true",
        choices=["16-true", "16-mixed", "bf16-true", "bf16-mixed", "32-true", "64-true", "64", "32", "16", "bf16"],
        help="precision",
    )
    parser.add_argument("--devices", type=int, default=-1, help="devices num")
    parser.add_argument("--nodes", type=int, default=1, help="nodes num")
    parser.add_argument(
        "--disable-benchmark", action="store_true", default=False, help="disable cudnn benchmark"
    )
    parser.add_argument(
        "--log-step", type=int, default=1, help="log training loss every n steps"
    )
    parser.add_argument(
        "--val-step", type=int, default=1600, help="valid and save every n steps, set 0 to valid and save every epoch"
    )

    opt = parser.parse_args()
    print(opt)

    if not os.path.exists("lightning_logs"):
        os.mkdir("lightning_logs")
    if not os.path.exists("sample"):
        os.mkdir("sample")
    pl.seed_everything(opt.seed)
    print("---load dataset---")
    if opt.config in config_name_list:
        config = MIDIModelConfig.from_name(opt.config)
    else:
        config = MIDIModelConfig.from_json_file(opt.config)
    tokenizer = config.tokenizer
    midi_list = get_midi_list(opt.data)
    random.shuffle(midi_list)
    full_dataset_len = len(midi_list)
    train_dataset_len = full_dataset_len - opt.data_val_split
    train_midi_list = midi_list[:train_dataset_len]
    val_midi_list = midi_list[train_dataset_len:]
    train_dataset = MidiDataset(train_midi_list, tokenizer, max_len=opt.max_len, aug=True, check_quality=opt.quality,
                                rand_start=True)
    val_dataset = MidiDataset(val_midi_list, tokenizer, max_len=opt.max_len, aug=False, check_quality=opt.quality,
                              rand_start=False)
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=opt.batch_size_train,
        shuffle=True,
        persistent_workers=True,
        num_workers=opt.workers_train,
        pin_memory=True,
        collate_fn=train_dataset.collate_fn
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=opt.batch_size_val,
        shuffle=False,
        persistent_workers=True,
        num_workers=opt.workers_val,
        pin_memory=True,
        collate_fn=val_dataset.collate_fn
    )
    print(f"train: {len(train_dataset)}  val: {len(val_dataset)}")
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    model = TrainMIDIModel(config, lr=opt.lr, weight_decay=opt.weight_decay,
                           warmup=opt.warmup_step, max_step=opt.max_step,
                           sample_seq=opt.sample_seq, gen_example_interval=opt.gen_example_interval,
                           example_batch=opt.batch_size_gen_example)
    if opt.ckpt:
        ckpt = torch.load(opt.ckpt, map_location="cpu")
        state_dict = ckpt.get("state_dict", ckpt)
        model.load_state_dict(state_dict, strict=False)
    elif opt.task == "lora":
        raise ValueError("--ckpt must be set to train lora")
    if opt.task == "lora":
        model.requires_grad_(False)
        lora_config = LoraConfig(
            r=64,
            target_modules=["q_proj", "o_proj", "k_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
            task_type=TaskType.CAUSAL_LM,
            bias="none",
            lora_alpha=128,
            lora_dropout=0
        )
        model.add_adapter(lora_config)
    print("---start train---")
    checkpoint_callback = ModelCheckpoint(
        monitor="val/loss",
        mode="min",
        save_top_k=1,
        save_last=True,
        auto_insert_metric_name=False,
        filename="epoch={epoch},loss={val/loss:.4f}",
    )
    callbacks = [checkpoint_callback]

    trainer = Trainer(
        precision=opt.precision,
        accumulate_grad_batches=opt.acc_grad,
        gradient_clip_val=opt.grad_clip,
        accelerator=opt.accelerator,
        devices=opt.devices,
        num_nodes=opt.nodes,
        max_steps=opt.max_step,
        benchmark=not opt.disable_benchmark,
        val_check_interval=opt.val_step or None,
        log_every_n_steps=1,
        strategy="auto",
        callbacks=callbacks,
    )
    ckpt_path = opt.resume
    if ckpt_path == "":
        ckpt_path = None
    print("---start train---")
    trainer.fit(model, train_dataloader, val_dataloader, ckpt_path=ckpt_path)