ML|Hung-yi Lee

import os
import json
import math
import csv
import random
from pathlib import Path
from tqdm import tqdm

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW, Optimizer
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import Dataset, DataLoader, random_split
from torch.nn.utils.rnn import pad_sequence


class myDataset(Dataset):
    """自定义数据集类，用于加载训练和验证数据。"""
    def __init__(self, data_dir, segment_len=128):
        self.data_dir = data_dir
        self.segment_len = segment_len

        # 加载说话人名称到ID的映射。
        mapping_path = Path(data_dir) / "mapping.json"
        mapping = json.load(mapping_path.open())
        self.speaker2id = mapping["speaker2id"]

        # 加载训练数据的元数据。
        metadata_path = Path(data_dir) / "metadata.json"
        metadata = json.load(open(metadata_path))["speakers"]

        # 获取说话人总数。
        self.speaker_num = len(metadata.keys())
        self.data = []
        for speaker in metadata.keys():
            for utterances in metadata[speaker]:
                self.data.append([utterances["feature_path"], self.speaker2id[speaker]])

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

    def __getitem__(self, index):
        feat_path, speaker = self.data[index]
        # 加载预处理的梅尔频谱图。
        mel = torch.load(os.path.join(self.data_dir, feat_path))

        # 将梅尔频谱图分割成固定长度的片段。
        if len(mel) > self.segment_len:
            # 随机选择片段的起始点。
            start = random.randint(0, len(mel) - self.segment_len)
            # 获取一个长度为segment_len的片段。
            mel = torch.FloatTensor(mel[start:start+self.segment_len])
        else:
            mel = torch.FloatTensor(mel)
        # 将说话人ID转换为long类型，方便后续计算损失。
        speaker = torch.FloatTensor([speaker]).long()
        return mel, speaker

    def get_speaker_number(self):
        return self.speaker_num


def collate_batch(batch):
    """整理一个批次的数据。"""
    mel, speaker = zip(*batch)
    # 对同一批次中的特征进行填充，使其长度相同。
    mel = pad_sequence(mel, batch_first=True, padding_value=-20)  # 填充值为很小的值。
    # mel: (batch size, length, 40)
    return mel, torch.FloatTensor(speaker).long()


def get_dataloader(data_dir, batch_size, n_workers):
    """生成数据加载器。"""
    dataset = myDataset(data_dir)
    speaker_num = dataset.get_speaker_number()
    # 将数据集分割为训练集和验证集。
    trainlen = int(0.9 * len(dataset))
    lengths = [trainlen, len(dataset) - trainlen]
    trainset, validset = random_split(dataset, lengths)

    train_loader = DataLoader(
        trainset,
        batch_size=batch_size,
        shuffle=True,
        drop_last=True,
        num_workers=n_workers,
        pin_memory=True,
        collate_fn=collate_batch,
    )
    valid_loader = DataLoader(
        validset,
        batch_size=batch_size,
        num_workers=n_workers,
        drop_last=True,
        pin_memory=True,
        collate_fn=collate_batch,
    )

    return train_loader, valid_loader, speaker_num

# 定义前馈模块
class FeedForwardModule(nn.Module):
    """Conformer中的前馈模块。"""
    def __init__(self, d_model, dim_feedforward, dropout):
        super(FeedForwardModule, self).__init__()
        self.seq = nn.Sequential(
            nn.LayerNorm(d_model),
            nn.Linear(d_model, dim_feedforward),
            nn.SiLU(),  # 使用SiLU（Swish）激活函数
            nn.Dropout(dropout),
            nn.Linear(dim_feedforward, d_model),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.seq(x)


# 定义卷积模块
class ConvolutionModule(nn.Module):
    """Conformer中的卷积模块。"""
    def __init__(self, d_model, kernel_size=31, dropout=0.1):
        super(ConvolutionModule, self).__init__()
        self.layer_norm = nn.LayerNorm(d_model)
        self.pointwise_conv1 = nn.Conv1d(d_model, 2 * d_model, kernel_size=1, padding=0)
        self.glu = nn.GLU(dim=1)
        self.depthwise_conv = nn.Conv1d(d_model, d_model, kernel_size=kernel_size,
                                        padding=(kernel_size - 1) // 2, groups=d_model)
        self.batch_norm = nn.BatchNorm1d(d_model)
        self.activation = nn.SiLU()  # 使用SiLU（Swish）激活函数
        self.pointwise_conv2 = nn.Conv1d(d_model, d_model, kernel_size=1, padding=0)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: (seq_len, batch_size, d_model)
        x = x.permute(1, 0, 2)  # (batch_size, seq_len, d_model)
        x = self.layer_norm(x)
        x = x.transpose(1, 2)    # (batch_size, d_model, seq_len)
        x = self.pointwise_conv1(x)
        x = self.glu(x)          # (batch_size, d_model, seq_len)
        x = self.depthwise_conv(x)
        x = self.batch_norm(x)
        x = self.activation(x)
        x = self.pointwise_conv2(x)
        x = self.dropout(x)
        x = x.transpose(1, 2)    # (batch_size, seq_len, d_model)
        x = x.permute(1, 0, 2)   # (seq_len, batch_size, d_model)
        return x


# 定义Conformer块
class ConformerBlock(nn.Module):
    """结合前馈、注意力和卷积模块的Conformer块。"""
    def __init__(self, d_model, nhead, dim_feedforward=256, kernel_size=31, dropout=0.1):
        super(ConformerBlock, self).__init__()
        self.ffn1 = FeedForwardModule(d_model, dim_feedforward, dropout)
        self.mha = nn.MultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)
        self.conv_module = ConvolutionModule(d_model, kernel_size, dropout)
        self.ffn2 = FeedForwardModule(d_model, dim_feedforward, dropout)
        self.layer_norm = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # x: (seq_len, batch_size, d_model)
        residual = x
        x = residual + 0.5 * self.ffn1(x)
        residual = x
        x, _ = self.mha(x, x, x)
        x = residual + x
        x = x + self.conv_module(x)
        x = x + 0.5 * self.ffn2(x)
        x = self.layer_norm(x)
        return x


# 定义分类器模型
class Classifier(nn.Module):
    def __init__(self, d_model=80, n_spks=600, num_layers=2, dropout=0.1):
        super().__init__()
        # 将输入特征的维度从40投影到d_model。
        self.prenet = nn.Linear(40, d_model)

        # 使用多个Conformer块代替Transformer层
        self.conformer_blocks = nn.ModuleList([
            ConformerBlock(d_model=d_model, nhead=2, dim_feedforward=256, dropout=dropout)
            for _ in range(num_layers)
        ])

        # 将特征维度从d_model投影到说话人数量。
        self.pred_layer = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, n_spks),
        )

    def forward(self, mels):
        """
        args:
          mels: (batch size, length, 40)
        return:
          out: (batch size, n_spks)
        """
        # out: (batch size, length, d_model)
        out = self.prenet(mels)
        # out: (length, batch size, d_model)
        out = out.permute(1, 0, 2)

        # 通过Conformer块
        for block in self.conformer_blocks:
            out = block(out)

        # out: (batch size, length, d_model)
        out = out.transpose(0, 1)
        # 平均池化
        stats = out.mean(dim=1)

        # out: (batch, n_spks)
        out = self.pred_layer(stats)
        return out


def get_cosine_schedule_with_warmup(
    optimizer: Optimizer,
    num_warmup_steps: int,
    num_training_steps: int,
    num_cycles: float = 0.5,
    last_epoch: int = -1,
):
    """
    创建一个学习率调度器，在预热期间学习率从0线性增加到初始值，然后按照余弦函数下降。

    Args:
      optimizer (:class:`~torch.optim.Optimizer`):
        优化器。
      num_warmup_steps (:obj:`int`):
        预热阶段的步数。
      num_training_steps (:obj:`int`):
        总的训练步数。
      num_cycles (:obj:`float`, `optional`, defaults to 0.5):
        余弦调度中的周期数。
      last_epoch (:obj:`int`, `optional`, defaults to -1):
        上一个epoch的索引。

    Return:
      :obj:`torch.optim.lr_scheduler.LambdaLR` 学习率调度器。
    """

    def lr_lambda(current_step):
        # 预热阶段
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        # 衰减阶段
        progress = float(current_step - num_warmup_steps) / float(
            max(1, num_training_steps - num_warmup_steps)
        )
        return max(
            0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)


def model_fn(batch, model, criterion, device):
    """将一个批次的数据前向传播通过模型。"""

    mels, labels = batch
    mels = mels.to(device)
    labels = labels.to(device)

    outs = model(mels)

    loss = criterion(outs, labels)

    # 获取具有最高概率的说话人ID。
    preds = outs.argmax(1)
    # 计算准确率。
    accuracy = torch.mean((preds == labels).float())

    return loss, accuracy


def valid(dataloader, model, criterion, device):
    """在验证集上进行验证。"""

    model.eval()
    running_loss = 0.0
    running_accuracy = 0.0
    pbar = tqdm(total=len(dataloader.dataset), ncols=0, desc="Valid", unit=" uttr")

    for i, batch in enumerate(dataloader):
        with torch.no_grad():
            loss, accuracy = model_fn(batch, model, criterion, device)
            running_loss += loss.item()
            running_accuracy += accuracy.item()

        pbar.update(dataloader.batch_size)
        pbar.set_postfix(
            loss=f"{running_loss / (i+1):.2f}",
            accuracy=f"{running_accuracy / (i+1):.2f}",
        )

    pbar.close()
    model.train()

    return running_accuracy / len(dataloader)


def parse_args():
    """参数设置"""
    config = {
        "data_dir": "./Dataset",      # 数据集路径
        "save_path": "model.ckpt",    # 模型保存路径
        "batch_size": 32,             # 批次大小
        "n_workers": 8,               # 数据加载器的工作线程数
        "valid_steps": 2000,          # 每隔多少步进行一次验证
        "warmup_steps": 1000,         # 学习率预热步数
        "save_steps": 10000,          # 每隔多少步保存一次模型
        "total_steps": 70000,         # 总训练步数
        "model_path": "model.ckpt",   # 推理时加载的模型路径
        "output_path": "output.csv",  # 推理结果输出路径
    }

    return config


def main():
    """主函数。"""
    # 解析参数
    config = parse_args()
    data_dir = config["data_dir"]
    save_path = config["save_path"]
    batch_size = config["batch_size"]
    n_workers = config["n_workers"]
    valid_steps = config["valid_steps"]
    warmup_steps = config["warmup_steps"]
    total_steps = config["total_steps"]
    save_steps = config["save_steps"]
    model_path = config["model_path"]
    output_path = config["output_path"]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[Info]: 使用 {device} 作为计算设备。")

    # 获取数据加载器和说话人数
    train_loader, valid_loader, speaker_num = get_dataloader(data_dir, batch_size, n_workers)
    train_iterator = iter(train_loader)
    print(f"[Info]: 数据加载完成。", flush=True)

    # 初始化模型、损失函数、优化器和学习率调度器
    model = Classifier(n_spks=speaker_num).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = AdamW(model.parameters(), lr=1e-3)
    scheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)
    print(f"[Info]: 模型创建完成。", flush=True)

    best_accuracy = -1.0
    best_state_dict = None

    pbar = tqdm(total=valid_steps, ncols=0, desc="Train", unit=" step")

    for step in range(total_steps):
        # 获取数据
        try:
            batch = next(train_iterator)
        except StopIteration:
            train_iterator = iter(train_loader)
            batch = next(train_iterator)

        loss, accuracy = model_fn(batch, model, criterion, device)
        batch_loss = loss.item()
        batch_accuracy = accuracy.item()

        # 更新模型
        loss.backward()
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

        # 日志记录
        pbar.update()
        pbar.set_postfix(
            loss=f"{batch_loss:.2f}",
            accuracy=f"{batch_accuracy:.2f}",
            step=step + 1,
        )

        # 进行验证
        if (step + 1) % valid_steps == 0:
            pbar.close()

            valid_accuracy = valid(valid_loader, model, criterion, device)

            # 保存最好的模型
            if valid_accuracy > best_accuracy:
                best_accuracy = valid_accuracy
                best_state_dict = model.state_dict()

            pbar = tqdm(total=valid_steps, ncols=0, desc="Train", unit=" step")

        # 保存模型
        if (step + 1) % save_steps == 0 and best_state_dict is not None:
            torch.save(best_state_dict, save_path)
            pbar.write(f"步骤 {step + 1}，最佳模型已保存。（准确率={best_accuracy:.4f}）")

    pbar.close()

    # 最终保存模型
    if best_state_dict is not None:
        torch.save(best_state_dict, save_path)
        print(f"[Info]: 训练完成，最佳模型已保存。（准确率={best_accuracy:.4f}）")

    # 推理部分
    # 加载测试数据集
    class InferenceDataset(Dataset):
        def __init__(self, data_dir):
            testdata_path = Path(data_dir) / "testdata.json"
            metadata = json.load(testdata_path.open())
            self.data_dir = data_dir
            self.data = metadata["utterances"]

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

        def __getitem__(self, index):
            utterance = self.data[index]
            feat_path = utterance["feature_path"]
            mel = torch.load(os.path.join(self.data_dir, feat_path))

            return feat_path, mel

    def inference_collate_batch(batch):
        """整理一个批次的数据。"""
        feat_paths, mels = zip(*batch)

        return feat_paths, torch.stack(mels)

    # 进行推理
    print("[Info]: 开始推理...", flush=True)
    mapping_path = Path(data_dir) / "mapping.json"
    mapping = json.load(mapping_path.open())

    inference_dataset = InferenceDataset(data_dir)
    inference_dataloader = DataLoader(
        inference_dataset,
        batch_size=1,
        shuffle=False,
        drop_last=False,
        num_workers=n_workers,
        collate_fn=inference_collate_batch,
    )
    print(f"[Info]: 推理数据加载完成，共有 {len(inference_dataset)} 条数据。", flush=True)

    # 加载训练好的模型
    model = Classifier(n_spks=speaker_num).to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()
    print(f"[Info]: 模型加载完成。", flush=True)

    # 进行预测
    results = [["Id", "Category"]]
    for feat_paths, mels in tqdm(inference_dataloader):
        with torch.no_grad():
            mels = mels.to(device)
            outs = model(mels)
            preds = outs.argmax(1).cpu().numpy()
            for feat_path, pred in zip(feat_paths, preds):
                results.append([feat_path, mapping["id2speaker"][str(pred)]])

    # 保存结果到CSV文件
    with open(output_path, 'w', newline='') as csvfile:
        writer = csv.writer(csvfile)
        writer.writerows(results)
    print(f"[Info]: 推理完成，结果已保存到 {output_path}。")


if __name__ == "__main__":
    main()