from __future__ import annotations import argparse from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from audidata.io.crops import RandomCrop from audidata.transforms import Mono from audidata.collate.default import collate_fn from torch.utils.data import DataLoader import pytorch_lightning as pl from pytorch_lightning.callbacks import ModelCheckpoint from llm.llama import Llama, LlamaConfig from utils import LinearWarmUp, parse_yaml class MusicLLM(pl.LightningModule): def __init__(self, configs: dict): super().__init__() self.save_hyperparameters(configs) # Configs self.configs = configs self.sample_rate = configs["sample_rate"] self.max_seq_len = configs["max_caption_len"] self.lr = float(configs["train"]["lr"]) self.warm_up_steps = configs["train"]["warm_up_steps"] # Audio codec (预训练模型，不参与训练) self.codec = self._get_audio_codec(self.configs) self.codec.eval() # 设置为 eval 模式 for param in self.codec.parameters(): param.requires_grad = False # 冻结 codec 参数 # LLM vocab_size = configs["audio_codec"]["vocab_size"] self.llm = self._get_llm(vocab_size) def _get_audio_codec(self, configs: dict) -> nn.Module: name = configs["audio_codec"]["name"] if name == "trancodec_fsq": from audio_codec.trancodec_fsq import Trancodec return Trancodec(configs) def _get_llm(self, vocab_size: int) -> nn.Module: config = LlamaConfig( block_size=self.configs["llm"]["block_size"], vocab_size=vocab_size, n_layer=self.configs["llm"]["n_layer"], n_head=self.configs["llm"]["n_head"], n_embd=self.configs["llm"]["n_embd"] ) return Llama(config=config) def forward(self, audio_codes): return self.llm(audio_codes) def training_step(self, batch, batch_idx): audio, _ = self._process_batch(batch) with torch.no_grad(): # 确保 codec 不计算梯度 self.codec.eval() # 强制 eval 模式 audio_codes = self.codec.encode(audio).squeeze(dim=2) # (b, t, q) # Prepare input and target input_codes = audio_codes[:, :-1] # (b, t-1, q) target_codes = audio_codes[:, 1:].to(torch.long) # (b, t-1, q) # Forward logits = self(input_codes) # (b, t-1, q) # Loss loss = F.cross_entropy( logits.flatten(0, 1), target_codes.flatten(0, 1), reduction="mean" ) self.log("train_loss", loss, prog_bar=True) return loss def validation_step(self, batch, batch_idx): audio, _ = self._process_batch(batch) with torch.no_grad(): # 确保 codec 不计算梯度 self.codec.eval() audio_codes = self.codec.encode(audio).squeeze(dim=2) input_codes = audio_codes[:, :-1] target_codes = audio_codes[:, 1:].to(torch.long) with torch.no_grad(): logits = self(input_codes) loss = F.cross_entropy( logits.flatten(0, 1), target_codes.flatten(0, 1), reduction="mean" ) self.log("val_loss", loss, prog_bar=True) return loss def configure_optimizers(self): # 只优化 llm 参数，排除 codec optimizer = torch.optim.AdamW(self.llm.parameters(), lr=self.lr) if self.warm_up_steps: lr_lambda = LinearWarmUp(self.warm_up_steps) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) return [optimizer], [scheduler] return optimizer def _process_batch(self, batch): return batch["audio"], batch["dataset_name"] def get_dataset(configs: dict, split: str) -> torch.utils.data.Dataset: datasets_split = "{}_datasets".format(split) for name in configs[datasets_split].keys(): if name == "Mozart": from audidata.datasets import Mozart dataset = Mozart( root=configs[datasets_split][name]["root"], split=configs[datasets_split][name]["split"], sr=configs["sample_rate"], crop=RandomCrop(clip_duration=configs["clip_duration"], end_pad=configs["clip_duration"] - 0.1), transform=Mono(), ) return dataset elif name == "MAESTRO": from audidata.datasets import MAESTRO from audidata.transforms.midi import PianoRoll dataset = MAESTRO( root=configs[datasets_split][name]["root"], split=configs[datasets_split][name]["split"], sr=configs["sample_rate"], crop=RandomCrop(clip_duration=configs["clip_duration"], end_pad=configs["clip_duration"] - 0.1), transform=Mono(), load_target=True, extend_pedal=True, target_transform=PianoRoll(fps=configs["fps"], pitches_num=128), ) return dataset else: raise ValueError(name) def main(args): # Parse configs configs = parse_yaml(args.config) # Datasets train_dataset = get_dataset(configs, "train") val_dataset = get_dataset(configs, "test") # DataLoaders train_loader = DataLoader( train_dataset, batch_size=configs["train"]["batch_size_per_device"], shuffle=True, num_workers=configs["train"]["num_workers"], collate_fn=collate_fn, pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=configs["train"]["batch_size_per_device"], shuffle=False, num_workers=configs["train"]["num_workers"], collate_fn=collate_fn, pin_memory=True ) # Model model = MusicLLM(configs) # Checkpoint callback checkpoint_callback = ModelCheckpoint( dirpath=Path("./checkpoints", Path(__file__).stem, Path(args.config).stem), filename="step={step}", every_n_train_steps=configs["train"]["save_every_n_steps"], save_top_k=-1 ) # Trainer trainer = pl.Trainer( max_steps=configs["train"]["training_steps"], accelerator="gpu", devices=configs["train"]["device"], strategy="ddp", precision=32, callbacks=[checkpoint_callback], check_val_every_n_epoch=configs["train"]["test_every_n_epoch"], ) # Train trainer.fit(model, train_loader, val_loader) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, required=True, help="Path of config yaml.") args = parser.parse_args() main(args)