Upload training process.ipynb

training process.ipynb

@@ -0,0 +1,1598 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e8306d9f",
+   "metadata": {},
+   "source": [
+    "## 初始化"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "51338b4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import time\n",
+    "from typing import List, Union, Optional\n",
+    "import math\n",
+    "from types import SimpleNamespace\n",
+    "import random\n",
+    "import glob\n",
+    "from pathlib import Path\n",
+    "import pickle\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "import torch.optim as optim\n",
+    "from torch.utils.data import DataLoader, IterableDataset, Dataset\n",
+    "\n",
+    "from transformers.configuration_utils import PretrainedConfig\n",
+    "from transformers.modeling_utils import PreTrainedModel\n",
+    "from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS\n",
+    "from transformers.activations import ACT2FN\n",
+    "\n",
+    "from einops import rearrange, pack, unpack\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import soundfile\n",
+    "import audiomentations\n",
+    "\n",
+    "import numpy as np\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "from pl_utils import BaseModule, LearningRateConfig, TrainingConfig"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e15cad0e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pl_utils import init_before_training\n",
+    "\n",
+    "\n",
+    "init_before_training(\n",
+    "    matmul_precision=\"medium\",\n",
+    "    empty_cache=False,\n",
+    "    seed=42,\n",
+    ")\n",
+    "\n",
+    "num_workers = 28"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a828912f",
+   "metadata": {},
+   "source": [
+    "## 定义"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9592af7a",
+   "metadata": {},
+   "source": [
+    "### Utils 定义"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "84dd1eec",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def loudness_db2linear(db):\n",
+    "    return 10 ** (db / 20)\n",
+    "\n",
+    "\n",
+    "def loudness_linear2db(linear):\n",
+    "    return 20 * np.log10(linear)\n",
+    "\n",
+    "\n",
+    "def inference_one_with_model(\n",
+    "    model,\n",
+    "    mixed_wave,\n",
+    "    chunk_size=44100 * 8,\n",
+    "    overlap_size=44100 * 4,\n",
+    "    batch_size=16,\n",
+    "    gap_size=44100 * 1,\n",
+    "):\n",
+    "    \"\"\"\n",
+    "    输入一段 (C, wave_length) 音频张量，使用模型推理，输出 (num_stems, C, wave_length) 音频张量。\n",
+    "    \"\"\"\n",
+    "    # 淡入淡出 窗口\n",
+    "    fade_size = chunk_size // 10\n",
+    "    window = torch.ones(chunk_size - 2 * gap_size)\n",
+    "    window[:fade_size] = torch.linspace(0, 1, fade_size)\n",
+    "    window[-fade_size:] = torch.linspace(1, 0, fade_size)\n",
+    "    window = F.pad(window, (gap_size, gap_size), value=0.0)\n",
+    "    window = window.to(mixed_wave.device)\n",
+    "\n",
+    "    with torch.inference_mode():\n",
+    "        wave_length = mixed_wave.shape[-1]\n",
+    "\n",
+    "        if wave_length <= chunk_size:\n",
+    "            num_chunks = 1\n",
+    "        else:\n",
+    "            num_chunks = math.ceil((wave_length - chunk_size) / overlap_size) + 1\n",
+    "\n",
+    "        required_length = (num_chunks - 1) * overlap_size + chunk_size\n",
+    "        padded_wave = F.pad(\n",
+    "            mixed_wave,\n",
+    "            (0, required_length - wave_length),\n",
+    "            mode=\"constant\",\n",
+    "        )\n",
+    "\n",
+    "        unfolded_chunks = padded_wave.unfold(\n",
+    "            dimension=-1,\n",
+    "            size=chunk_size,\n",
+    "            step=overlap_size,\n",
+    "        )  # (C, num_chunks, chunk_size)\n",
+    "        batch = unfolded_chunks.permute(1, 0, 2)  # (num_chunks, C, chunk_size)\n",
+    "\n",
+    "        output_chunks = []\n",
+    "        for i in range(0, num_chunks, batch_size):\n",
+    "            chunk_batch = batch[i : i + batch_size]\n",
+    "            output_chunk = model(chunk_batch)\n",
+    "            output_chunks.append(output_chunk)\n",
+    "        batch = torch.cat(output_chunks, dim=0)  # (num_chunks, num_stems, C, chunk_size)\n",
+    "\n",
+    "        _, num_stems, C, _ = batch.shape\n",
+    "        batch = batch.view(num_chunks, -1, chunk_size).permute(1, 0, 2)  # (num_stems * C, num_chunks, chunk_size)\n",
+    "        batch = batch * window\n",
+    "        output_result_buffer = F.fold(\n",
+    "            batch.permute(0, 2, 1),\n",
+    "            output_size=(1, required_length),\n",
+    "            kernel_size=(1, chunk_size),\n",
+    "            stride=(1, overlap_size),\n",
+    "        )  # (num_stems * C, 1, 1, required_length)\n",
+    "\n",
+    "        window_for_fold = window.expand(1, 1, -1).repeat(1, num_chunks, 1)\n",
+    "        weighted_sum_counter = F.fold(\n",
+    "            window_for_fold.permute(0, 2, 1),\n",
+    "            output_size=(1, required_length),\n",
+    "            kernel_size=(1, chunk_size),\n",
+    "            stride=(1, overlap_size),\n",
+    "        )  # (1, 1, 1, required_length)\n",
+    "\n",
+    "        output_result_buffer = output_result_buffer.view(num_stems, C, -1)  # (num_stems, C, required_length)\n",
+    "        weighted_sum_counter = weighted_sum_counter.view(1, 1, -1)\n",
+    "        weighted_sum_counter.clamp_min_(1e-8)\n",
+    "\n",
+    "        final_output = (output_result_buffer / weighted_sum_counter)[:, :, :wave_length]\n",
+    "\n",
+    "    return final_output"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "68c460af",
+   "metadata": {},
+   "source": [
+    "### Dataset 定义"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71aaa349",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class AugmentDataset(IterableDataset):\n",
+    "    \"\"\"\n",
+    "    用于 MUSDB18HQ 数据的、含有数据增强的 Dataset。返回分块音频。\n",
+    "\n",
+    "    期望的数据目录结构：\n",
+    "\n",
+    "        dataset/\n",
+    "        ├── A Classic Education - NightOwl\n",
+    "        │   ├── bass.wav\n",
+    "        │   ├── drums.wav\n",
+    "        │   ├── mixture.wav\n",
+    "        │   ├── other.wav\n",
+    "        │   └── vocals.wav\n",
+    "        ├── Actions - Devil's Words\n",
+    "        │   ├── bass.wav\n",
+    "        │   ├── drums.wav\n",
+    "        │   ├── mixture.wav\n",
+    "        │   ├── other.wav\n",
+    "        │   └── vocals.wav\n",
+    "        ···\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        data_path,\n",
+    "        wave_chunk_size=44100 * 8,\n",
+    "        sample_rate=44100,\n",
+    "        same_stem_mixup_prob=[0.2, 0.02],\n",
+    "        same_stem_mixup_loudness_range=[-3, 3],\n",
+    "        stem_names=[\"bass\", \"drums\", \"other\", \"vocals\"],\n",
+    "        debug=False,\n",
+    "    ):\n",
+    "        if type(data_path) is not list:\n",
+    "            data_path = [data_path]\n",
+    "        self.data_path = [Path(p) for p in data_path]\n",
+    "\n",
+    "        self.wave_chunk_size = wave_chunk_size\n",
+    "        self.sample_rate = sample_rate\n",
+    "\n",
+    "        self.same_stem_mixup_prob = same_stem_mixup_prob\n",
+    "        self.same_stem_mixup_loudness_range = same_stem_mixup_loudness_range\n",
+    "        self.stem_names = stem_names\n",
+    "\n",
+    "        self.metadata = self._get_metadata()\n",
+    "\n",
+    "        self.augments = audiomentations.Compose(\n",
+    "            [\n",
+    "                # 极性反转\n",
+    "                audiomentations.PolarityInversion(p=0.5),\n",
+    "                # 音高偏移\n",
+    "                # audiomentations.PitchShift(\n",
+    "                #     min_semitones=-5,\n",
+    "                #     max_semitones=5,\n",
+    "                #     p=0.5,\n",
+    "                # ),\n",
+    "                # 七频段 eq 随机调整\n",
+    "                audiomentations.SevenBandParametricEQ(\n",
+    "                    min_gain_db=-9,\n",
+    "                    max_gain_db=9,\n",
+    "                    p=1.0,\n",
+    "                ),\n",
+    "                # tanh 失真\n",
+    "                audiomentations.TanhDistortion(\n",
+    "                    min_distortion=0.1,\n",
+    "                    max_distortion=0.6,\n",
+    "                    p=0.5,\n",
+    "                ),\n",
+    "                # 低品质失真\n",
+    "                audiomentations.Mp3Compression(\n",
+    "                    min_bitrate=32,\n",
+    "                    max_bitrate=256,\n",
+    "                    p=0.4,\n",
+    "                ),\n",
+    "                # 拉伸\n",
+    "                # audiomentations.TimeStretch(\n",
+    "                #     min_rate=0.8,\n",
+    "                #     max_rate=1.25,\n",
+    "                #     p=1.0,\n",
+    "                # ),\n",
+    "                # 随机音量\n",
+    "                # audiomentations.GainTransition(\n",
+    "                #     min_gain_db=-3,\n",
+    "                #     max_gain_db=3,\n",
+    "                #     min_duration=0.5,\n",
+    "                #     max_duration=4.0,\n",
+    "                #     p=1.0,\n",
+    "                # ),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "        self.file_handles = {}\n",
+    "        self.debug = debug\n",
+    "\n",
+    "    def _get_one_of_metadata(self, data_path):\n",
+    "        song_paths = [p for p in data_path.iterdir() if p.is_dir()]\n",
+    "        # 读取缓存\n",
+    "        cache_path = data_path / \"metadata.pkl\"\n",
+    "        if cache_path.exists():\n",
+    "            with open(cache_path, \"rb\") as f:\n",
+    "                song_metadata = pickle.load(f)\n",
+    "                cache_paths = [m[0] for m in song_metadata]\n",
+    "                # 文件没有改动，直接使用缓存\n",
+    "                if set(cache_paths) == set(song_paths):\n",
+    "                    return song_metadata\n",
+    "\n",
+    "        # 构建缓存\n",
+    "        song_metadata = []\n",
+    "        for song_path in tqdm(song_paths, desc=\"Scanning dataset\"):\n",
+    "            wave_files = [f for f in song_path.iterdir() if f.is_file() and f.stem in self.stem_names]\n",
+    "\n",
+    "            lengths = []\n",
+    "            for wave_file in wave_files:\n",
+    "                data, samplerate = soundfile.read(wave_file)\n",
+    "                assert samplerate == self.sample_rate, f\"Sample rate {samplerate} is not desired {self.sample_rate}\"\n",
+    "                track_length = len(data)\n",
+    "                lengths.append(track_length)\n",
+    "            if len(set(lengths)) > 1:\n",
+    "                print(f\"Warning: Inconsistent track lengths found in {song_path}. Using min length: {min(lengths)}\")\n",
+    "\n",
+    "            stem_file_dict = {f.stem: f for f in wave_files}\n",
+    "            song_metadata.append((song_path, min(lengths), stem_file_dict))\n",
+    "\n",
+    "        # 保存缓存\n",
+    "        with open(cache_path, \"wb\") as f:\n",
+    "            pickle.dump(song_metadata, f)\n",
+    "\n",
+    "        return song_metadata\n",
+    "\n",
+    "    def _get_metadata(self):\n",
+    "        all_metadata = []\n",
+    "        for p in self.data_path:\n",
+    "            metadata = self._get_one_of_metadata(p)\n",
+    "            all_metadata.extend(metadata)\n",
+    "        return all_metadata\n",
+    "\n",
+    "    def _load_random_wave(self, stem_name):\n",
+    "        \"\"\"\n",
+    "        从 self.metadata 选取出指定 stem_name 的音轨。来源歌曲、截取位置都随机。\n",
+    "\n",
+    "        截取长度由 `self.wave_chunk_size` 决定。\n",
+    "        \"\"\"\n",
+    "\n",
+    "        # 尝试 10 次，保证音频响度大于 -50dB\n",
+    "        for _ in range(10):\n",
+    "            song_path, length, stem_file_dict = random.choice(self.metadata)\n",
+    "\n",
+    "            # random offset within track\n",
+    "            offset = np.random.randint(length - self.wave_chunk_size + 1)\n",
+    "            # get or open cached file handle\n",
+    "            file_path = stem_file_dict[stem_name]\n",
+    "            if file_path not in self.file_handles:\n",
+    "                self.file_handles[file_path] = soundfile.SoundFile(str(file_path), mode='r')\n",
+    "            handle = self.file_handles[file_path]\n",
+    "            # seek and read chunk\n",
+    "            handle.seek(offset)\n",
+    "            wave = handle.read(self.wave_chunk_size, dtype='float32')\n",
+    "            wave = wave.T  # (channel, time)\n",
+    "            if len(wave.shape) == 1:  # 对 mono 音频添加 channel 维度\n",
+    "                wave = np.expand_dims(wave, axis=0)\n",
+    "\n",
+    "            rms = np.sqrt(np.mean(wave**2))\n",
+    "            if rms > loudness_db2linear(-50):\n",
+    "                break\n",
+    "\n",
+    "            if self.debug:\n",
+    "                print(f\"Warning: sampled very silent audio from {file_path} (rms={rms:.6f})\")\n",
+    "        # augmentation\n",
+    "        wave = self._apply_augment(wave, stem_name)\n",
+    "\n",
+    "        return wave\n",
+    "\n",
+    "    def _load_random_stems(self):\n",
+    "        \"\"\"\n",
+    "        加载随机的 self.stem_names 分轨。\n",
+    "\n",
+    "        包含的数据增强：\n",
+    "\n",
+    "        - 单个 stem 的来源歌曲和截取位置都随机（由 `self._load_random_track()` 实现）\n",
+    "        - 单个 stem 可能是多个同类型 stem 混合获得，概率由 `self.same_stem_mixup_prob` 决定\n",
+    "            - 混合 stem 时各个 stem 的响度在 `self.same_stem_mixup_loudness_range` 范围内随机\n",
+    "        \"\"\"\n",
+    "        waves = []\n",
+    "        for stem_name in self.stem_names:\n",
+    "            wave = self._load_random_wave(stem_name)\n",
+    "\n",
+    "            mixup_waves = [wave]\n",
+    "            for prob in self.same_stem_mixup_prob:\n",
+    "                if random.uniform(0, 1) < prob:\n",
+    "                    wave2 = self._load_random_wave(stem_name)\n",
+    "                    mixup_waves.append(wave2)\n",
+    "\n",
+    "            mixup_waves = np.stack(mixup_waves, axis=0)\n",
+    "\n",
+    "            # 在 self.same_stem_mixup_loudness_range 范围内的随机响度\n",
+    "            loudness = np.random.uniform(\n",
+    "                low=loudness_db2linear(self.same_stem_mixup_loudness_range[0]),\n",
+    "                high=loudness_db2linear(self.same_stem_mixup_loudness_range[1]),\n",
+    "                size=(len(mixup_waves),),\n",
+    "            )\n",
+    "            mixup_waves *= loudness[:, None, None]\n",
+    "            mixup_wave = mixup_waves.mean(axis=0)\n",
+    "\n",
+    "            waves.append(mixup_wave)\n",
+    "\n",
+    "        waves = np.stack(waves, axis=0)\n",
+    "\n",
+    "        return waves\n",
+    "\n",
+    "    def _apply_augment(self, wave, stem_name):\n",
+    "        # Channel shuffle\n",
+    "        if random.uniform(0, 1) < 0.5:\n",
+    "            wave = wave[::-1].copy()\n",
+    "\n",
+    "        # self.stem_augment\n",
+    "        wave = self.augments(samples=wave, sample_rate=self.sample_rate)\n",
+    "\n",
+    "        return wave\n",
+    "\n",
+    "    def __iter__(self):\n",
+    "        while True:\n",
+    "            waves = self._load_random_stems()\n",
+    "\n",
+    "            # 随机分轨音量\n",
+    "            loudnesses = np.random.uniform(\n",
+    "                low=loudness_db2linear(-3),\n",
+    "                high=loudness_db2linear(3),\n",
+    "                size=(len(waves),),\n",
+    "            )\n",
+    "            # 各个 stem 有 10% 概率变为空音频\n",
+    "            loudnesses *= (np.random.uniform(0, 1, size=(len(waves),)) > 0.1).astype(np.float32)\n",
+    "            # 施加到 waves 上\n",
+    "            waves *= loudnesses[:, None, None]\n",
+    "\n",
+    "            # 获得混合音频\n",
+    "            mixed_wave = waves.sum(0)\n",
+    "\n",
+    "            yield waves, mixed_wave\n",
+    "\n",
+    "    def __del__(self):\n",
+    "        # Close any open SoundFile handles when dataset is destroyed\n",
+    "        for handle in self.file_handles.values():\n",
+    "            try:\n",
+    "                handle.close()\n",
+    "            except Exception:\n",
+    "                pass\n",
+    "\n",
+    "\n",
+    "class ValidationDataset(Dataset):\n",
+    "    \"\"\"\n",
+    "    用于 MUSDB18HQ 数据的、用于验证的 Dataset。返回完整音频。\n",
+    "\n",
+    "    期望的数据目录结构：\n",
+    "\n",
+    "        dataset/\n",
+    "        ├── A Classic Education - NightOwl\n",
+    "        │   ├── bass.wav\n",
+    "        │   ├── drums.wav\n",
+    "        │   ├── mixture.wav\n",
+    "        │   ├── other.wav\n",
+    "        │   └── vocals.wav\n",
+    "        ├── Actions - Devil's Words\n",
+    "        │   ├── bass.wav\n",
+    "        │   ├── drums.wav\n",
+    "        │   ├── mixture.wav\n",
+    "        │   ├── other.wav\n",
+    "        │   └── vocals.wav\n",
+    "        ···\n",
+    "    \"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        data_path,\n",
+    "        sample_rate=44100,\n",
+    "        stem_names=[\"bass\", \"drums\", \"other\", \"vocals\"],\n",
+    "    ):\n",
+    "        self.data_path = Path(data_path)\n",
+    "        self.sample_rate = sample_rate\n",
+    "        self.stem_names = stem_names\n",
+    "\n",
+    "        self.metadata = self._get_metadata()\n",
+    "\n",
+    "    def _get_metadata(self):\n",
+    "        song_paths = [p for p in self.data_path.iterdir() if p.is_dir()]\n",
+    "        # 读取缓存\n",
+    "        cache_path = self.data_path / \"metadata.pkl\"\n",
+    "        if cache_path.exists():\n",
+    "            with open(cache_path, \"rb\") as f:\n",
+    "                song_metadata = pickle.load(f)\n",
+    "                cache_paths = [m[0] for m in song_metadata]\n",
+    "                # 文件没有改动，直接使用缓存\n",
+    "                if set(cache_paths) == set(song_paths):\n",
+    "                    return song_metadata\n",
+    "\n",
+    "        # 构建缓存\n",
+    "        song_metadata = []\n",
+    "        for song_path in tqdm(song_paths, desc=\"Scanning dataset\"):\n",
+    "            wave_files = [f for f in song_path.iterdir() if f.is_file() and f.stem in self.stem_names]\n",
+    "\n",
+    "            lengths = []\n",
+    "            for wave_file in wave_files:\n",
+    "                data, samplerate = soundfile.read(wave_file)\n",
+    "                assert samplerate == self.sample_rate, f\"Sample rate {samplerate} is not desired {self.sample_rate}\"\n",
+    "                track_length = len(data)\n",
+    "                lengths.append(track_length)\n",
+    "            if len(set(lengths)) > 1:\n",
+    "                print(f\"Warning: Inconsistent track lengths found in {song_path}. Using min length: {min(lengths)}\")\n",
+    "\n",
+    "            stem_file_dict = {f.stem: f for f in wave_files}\n",
+    "            song_metadata.append((song_path, min(lengths), stem_file_dict))\n",
+    "\n",
+    "        # 保存缓存\n",
+    "        with open(cache_path, \"wb\") as f:\n",
+    "            pickle.dump(song_metadata, f)\n",
+    "\n",
+    "        return song_metadata\n",
+    "\n",
+    "    def __len__(self):\n",
+    "        return len(self.metadata)\n",
+    "\n",
+    "    def __getitem__(self, index):\n",
+    "        song_path, length, stem_file_dict = self.metadata[index]\n",
+    "\n",
+    "        waves = []\n",
+    "        for stem_name in self.stem_names:\n",
+    "            stem_file = stem_file_dict[stem_name]\n",
+    "            wave = soundfile.read(\n",
+    "                stem_file,\n",
+    "                dtype=\"float32\",\n",
+    "            )[0]\n",
+    "            wave = wave.T\n",
+    "            if len(wave.shape) == 1:  # 对 mono 音频添加 channel 维度\n",
+    "                wave = np.expand_dims(wave, axis=0)\n",
+    "            waves.append(wave)\n",
+    "\n",
+    "        waves = np.stack(waves, axis=0)  # (stem, channel, time)\n",
+    "\n",
+    "        # 获得混合音频\n",
+    "        mixed_wave = waves.sum(0)\n",
+    "\n",
+    "        return waves, mixed_wave"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22caec1a",
+   "metadata": {},
+   "source": [
+    "### ModuleConfig 定义"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "591a48cd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "DEFAULT_FREQS_PER_BANDS = (\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    2,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    4,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    12,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    24,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    48,\n",
+    "    128,\n",
+    "    129,\n",
+    ")\n",
+    "\n",
+    "\n",
+    "class BSRoformerConfig(PretrainedConfig):\n",
+    "\n",
+    "    model_type = \"bs_roformer\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        hidden_size=384,\n",
+    "        depth=6,\n",
+    "        num_input_channel=1,\n",
+    "        num_stems=1,\n",
+    "        time_transformer_depth=2,\n",
+    "        freq_transformer_depth=2,\n",
+    "        freqs_per_bands: tuple[int, ...] = DEFAULT_FREQS_PER_BANDS,\n",
+    "        attention_dropout=0.0,\n",
+    "        num_attention_heads=8,\n",
+    "        num_key_value_heads=8,\n",
+    "        intermediate_size=384 * 4,\n",
+    "        #\n",
+    "        stft_n_fft=2048,\n",
+    "        stft_hop_length=512,\n",
+    "        stft_win_length=2048,\n",
+    "        mask_estimator_depth=2,\n",
+    "        multi_stft_loss_weight=1.0,  # TODO 权重降低会发生什么\n",
+    "        multi_stft_loss_window_sizes: tuple[int, ...] = (4096, 2048, 1024, 512, 256),\n",
+    "        multi_stft_loss_hop_size=147,\n",
+    "        rms_norm_eps=1e-6,\n",
+    "        rope_theta=10000.0,\n",
+    "        #\n",
+    "        initializer_range=0.02,\n",
+    "        register_token_num=4,\n",
+    "        **kwargs,\n",
+    "    ):\n",
+    "        self.hidden_size = hidden_size\n",
+    "        self.depth = depth\n",
+    "        self.num_input_channel = num_input_channel\n",
+    "        self.num_stems = num_stems\n",
+    "        self.time_transformer_depth = time_transformer_depth\n",
+    "        self.freq_transformer_depth = freq_transformer_depth\n",
+    "        self.freqs_per_bands = freqs_per_bands\n",
+    "        self.attention_dropout = attention_dropout\n",
+    "        self.num_attention_heads = num_attention_heads\n",
+    "        self.num_key_value_heads = num_key_value_heads\n",
+    "        self.intermediate_size = intermediate_size\n",
+    "\n",
+    "        self.stft_n_fft = stft_n_fft\n",
+    "        self.stft_hop_length = stft_hop_length\n",
+    "        self.stft_win_length = stft_win_length\n",
+    "\n",
+    "        self.mask_estimator_depth = mask_estimator_depth\n",
+    "        self.multi_stft_loss_weight = multi_stft_loss_weight\n",
+    "        self.multi_stft_loss_window_sizes = multi_stft_loss_window_sizes\n",
+    "        self.multi_stft_loss_hop_size = multi_stft_loss_hop_size\n",
+    "        self.rms_norm_eps = rms_norm_eps\n",
+    "        self.rope_theta = rope_theta\n",
+    "\n",
+    "        self.initializer_range = initializer_range\n",
+    "        self.register_token_num = register_token_num\n",
+    "\n",
+    "        super().__init__(**kwargs)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ba4ce953",
+   "metadata": {},
+   "source": [
+    "### 模型定义"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "48b33373",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# RoPE\n",
+    "class BSRoformerRotaryEmbedding(nn.Module):\n",
+    "    def __init__(self, dim, theta=10000.0):\n",
+    "        super().__init__()\n",
+    "        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))\n",
+    "        self.register_buffer(\"inv_freq\", inv_freq)\n",
+    "\n",
+    "    def forward(self, x, seq_len: int):\n",
+    "        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)\n",
+    "        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n",
+    "        emb = torch.cat((freqs, freqs), dim=-1)\n",
+    "        return emb.cos(), emb.sin()\n",
+    "\n",
+    "\n",
+    "def rotate_half(x):\n",
+    "    x1 = x[..., : x.shape[-1] // 2]\n",
+    "    x2 = x[..., x.shape[-1] // 2 :]\n",
+    "    return torch.cat((-x2, x1), dim=-1)\n",
+    "\n",
+    "\n",
+    "def apply_rotary_pos_emb(q, k, cos, sin):\n",
+    "    q_embed = (q * cos) + (rotate_half(q) * sin)\n",
+    "    k_embed = (k * cos) + (rotate_half(k) * sin)\n",
+    "    return q_embed, k_embed\n",
+    "\n",
+    "\n",
+    "class RotaryEmbedding(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "        self.head_dim = config.hidden_size // config.num_attention_heads\n",
+    "        inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim))\n",
+    "        self.register_buffer(\"inv_freq\", inv_freq)\n",
+    "\n",
+    "    def forward(self, x, position_ids):\n",
+    "        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)\n",
+    "        position_ids_expanded = position_ids[:, None, :].float()\n",
+    "\n",
+    "        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != \"mps\" else \"cpu\"\n",
+    "        with torch.autocast(device_type=device_type, enabled=False):  # Force float32\n",
+    "            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)\n",
+    "            emb = torch.cat((freqs, freqs), dim=-1)\n",
+    "            cos = emb.cos()\n",
+    "            sin = emb.sin()\n",
+    "\n",
+    "        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)\n",
+    "\n",
+    "\n",
+    "# Attention\n",
+    "class BSRoformerMLP(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "        self.config = config\n",
+    "        self.hidden_size = config.hidden_size\n",
+    "        self.intermediate_size = config.intermediate_size\n",
+    "        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)\n",
+    "        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)\n",
+    "        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)\n",
+    "        self.act_fn = ACT2FN[\"gelu\"]\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))\n",
+    "        return down_proj\n",
+    "\n",
+    "\n",
+    "class BSRoformerAttention(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "        self.is_causal = False\n",
+    "        self.config = config\n",
+    "\n",
+    "        self.head_dim = config.hidden_size // config.num_attention_heads\n",
+    "        self.scaling = self.head_dim**-0.5\n",
+    "        self.attention_dropout = config.attention_dropout\n",
+    "\n",
+    "        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads\n",
+    "\n",
+    "        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)\n",
+    "        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)\n",
+    "        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)\n",
+    "        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        hidden_states,\n",
+    "        position_embeddings: tuple[torch.Tensor, torch.Tensor],\n",
+    "        attention_mask=None,\n",
+    "    ):\n",
+    "        input_shape = hidden_states.size()[:-1]\n",
+    "        hidden_shape = (*input_shape, -1, self.head_dim)  # b, n, d -> b, n, -1, d'\n",
+    "\n",
+    "        # proj\n",
+    "        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n",
+    "        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n",
+    "        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n",
+    "\n",
+    "        # positional embeddings\n",
+    "        cos, sin = position_embeddings\n",
+    "        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n",
+    "\n",
+    "        # multi-group attention\n",
+    "        # key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)\n",
+    "        # value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)\n",
+    "\n",
+    "        attention_interface = ALL_ATTENTION_FUNCTIONS[\"sdpa\"]\n",
+    "\n",
+    "        attn_output, attn_weights = attention_interface(\n",
+    "            self,\n",
+    "            query_states,\n",
+    "            key_states,\n",
+    "            value_states,\n",
+    "            attention_mask,\n",
+    "            dropout=0.0 if not self.training else self.attention_dropout,\n",
+    "            scaling=self.scaling,\n",
+    "        )\n",
+    "\n",
+    "        attn_output = attn_output.reshape(*input_shape, -1).contiguous()\n",
+    "        attn_output = self.o_proj(attn_output)\n",
+    "\n",
+    "        return attn_output, attn_weights\n",
+    "\n",
+    "\n",
+    "class BSRoformerLayer(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "        self.self_attn = BSRoformerAttention(config)\n",
+    "        self.mlp = BSRoformerMLP(config)\n",
+    "\n",
+    "        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "        self.post_attention_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        hidden_states,\n",
+    "        position_embeddings,\n",
+    "        attention_mask,\n",
+    "    ):\n",
+    "        # Self Attention\n",
+    "        residual = hidden_states\n",
+    "        hidden_states = self.input_layernorm(hidden_states)\n",
+    "        hidden_states, _ = self.self_attn(\n",
+    "            hidden_states,\n",
+    "            position_embeddings,\n",
+    "            attention_mask,\n",
+    "        )\n",
+    "        hidden_states = hidden_states + residual\n",
+    "\n",
+    "        # Fully Connected\n",
+    "        residual = hidden_states\n",
+    "        hidden_states = self.post_attention_layernorm(hidden_states)\n",
+    "        hidden_states = self.mlp(hidden_states)\n",
+    "        hidden_states = hidden_states + residual\n",
+    "\n",
+    "        return hidden_states\n",
+    "\n",
+    "\n",
+    "class BSRoformerAxialTransformer(nn.Module):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        config: BSRoformerConfig,\n",
+    "        transformer_depth: int,\n",
+    "        is_time_transformer: bool,\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        self.layers = nn.ModuleList([BSRoformerLayer(config) for _ in range(transformer_depth)])\n",
+    "        self.is_time_transformer = is_time_transformer\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        hidden_states,\n",
+    "        position_embeddings,\n",
+    "        attention_mask,\n",
+    "    ):\n",
+    "        if self.is_time_transformer:\n",
+    "            hidden_states = rearrange(hidden_states, 'b t f d -> b f t d')\n",
+    "\n",
+    "        # merge batch\n",
+    "        b, seq_len_1, seq_len_2, d = hidden_states.shape\n",
+    "        hidden_states = rearrange(hidden_states, 'b n m d -> (b n) m d')\n",
+    "\n",
+    "        for layer in self.layers:\n",
+    "            hidden_states = layer(\n",
+    "                hidden_states,\n",
+    "                position_embeddings,\n",
+    "                attention_mask,\n",
+    "            )\n",
+    "\n",
+    "        # unpack batch\n",
+    "        hidden_states = rearrange(hidden_states, '(b n) m d -> b n m d', b=b)\n",
+    "\n",
+    "        if self.is_time_transformer:\n",
+    "            hidden_states = rearrange(hidden_states, 'b f t d -> b t f d')\n",
+    "\n",
+    "        return hidden_states\n",
+    "\n",
+    "\n",
+    "# BandSplit & MaskEstimator\n",
+    "class BandSplit(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "        self.dim_inputs = tuple(2 * f * config.num_input_channel for f in config.freqs_per_bands)\n",
+    "        self.to_features = nn.ModuleList(\n",
+    "            [\n",
+    "                nn.Sequential(nn.RMSNorm(dim_in, eps=config.rms_norm_eps), nn.Linear(dim_in, config.hidden_size))\n",
+    "                for dim_in in self.dim_inputs\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_split = x.split(self.dim_inputs, dim=-1)\n",
+    "        outs = [to_feature(split_input) for split_input, to_feature in zip(x_split, self.to_features)]\n",
+    "        return torch.stack(outs, dim=-2)\n",
+    "\n",
+    "\n",
+    "def MLP(dim_in, dim_out, dim_hidden, depth, activation=nn.Tanh):\n",
+    "    net = []\n",
+    "    dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out)\n",
+    "    for i, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])):\n",
+    "        net.append(nn.Linear(layer_dim_in, layer_dim_out))\n",
+    "        if i < len(dims) - 2:\n",
+    "            net.append(activation())\n",
+    "    return nn.Sequential(*net)\n",
+    "\n",
+    "\n",
+    "class MaskEstimator(nn.Module):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__()\n",
+    "\n",
+    "        class MiniGeGLU(nn.Module):\n",
+    "\n",
+    "            def __init__(self, out_size):\n",
+    "                super().__init__()\n",
+    "                self.gate_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)\n",
+    "                self.up_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=False)\n",
+    "                self.down_proj = nn.Linear(config.hidden_size, out_size, bias=False)\n",
+    "                self.act_fn = nn.GELU()\n",
+    "\n",
+    "            def forward(self, x):\n",
+    "                down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))\n",
+    "                return down_proj\n",
+    "\n",
+    "        dim_inputs = tuple(2 * f * config.num_input_channel for f in config.freqs_per_bands)\n",
+    "        # self.to_freq_mlps = nn.ModuleList([MiniGeGLU(dim_in) for dim_in in dim_inputs])\n",
+    "        self.to_freq_mlps = nn.ModuleList([nn.Linear(config.hidden_size, dim_in) for dim_in in dim_inputs])\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        x_unbind = x.unbind(dim=-2)\n",
+    "        outs = [mlp(band_features) for band_features, mlp in zip(x_unbind, self.to_freq_mlps)]\n",
+    "        return torch.cat(outs, dim=-1)\n",
+    "\n",
+    "\n",
+    "# Main Model\n",
+    "class BSRoformerPreTrainedModel(PreTrainedModel):\n",
+    "    config_class = BSRoformerConfig\n",
+    "    base_model_prefix = \"model\"\n",
+    "    supports_gradient_checkpointing = True\n",
+    "    _no_split_modules = [\"BSRoformerLayer\"]\n",
+    "\n",
+    "\n",
+    "class BSRoformerModel(BSRoformerPreTrainedModel):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__(config)\n",
+    "        self.config = config\n",
+    "        self.band_split = BandSplit(config)\n",
+    "        self.layers = nn.ModuleList(\n",
+    "            nn.ModuleList(\n",
+    "                [\n",
+    "                    BSRoformerAxialTransformer(config, config.time_transformer_depth, is_time_transformer=True),\n",
+    "                    BSRoformerAxialTransformer(config, config.freq_transformer_depth, is_time_transformer=False),\n",
+    "                ]\n",
+    "            )\n",
+    "            for _ in range(config.depth)\n",
+    "        )\n",
+    "        self.rotary_emb = RotaryEmbedding(config)\n",
+    "        self.final_norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)\n",
+    "\n",
+    "        rn = config.register_token_num\n",
+    "        self.register_tokens = nn.Parameter(torch.normal(0, 0.02, size=(rn, rn, config.hidden_size)))\n",
+    "\n",
+    "        self.post_init()\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        x,\n",
+    "        position_ids=None,\n",
+    "    ):\n",
+    "        hidden_states = self.band_split(x)\n",
+    "\n",
+    "        b, t, n, h = hidden_states.shape  # [batch, t, n, hidden_size]\n",
+    "\n",
+    "        if position_ids is None:\n",
+    "            position_ids = torch.arange(t, device=hidden_states.device).unsqueeze(0)\n",
+    "        pos_embeds = self.rotary_emb(hidden_states, position_ids)\n",
+    "        pos_embeds_for_freq = self.rotary_emb(\n",
+    "            hidden_states,\n",
+    "            torch.arange(n, device=hidden_states.device).unsqueeze(0),\n",
+    "        )\n",
+    "\n",
+    "        # add register tokens\n",
+    "        rn = self.config.register_token_num\n",
+    "        hidden_states = F.pad(hidden_states, (0, 0, 0, rn, 0, rn))\n",
+    "        hidden_states[:, t:, n:, :] = self.register_tokens\n",
+    "\n",
+    "        def pad_rope(cos, sin):\n",
+    "            cos_padded = F.pad(cos, (0, 0, 0, rn), value=1.0)\n",
+    "            sin_padded = F.pad(sin, (0, 0, 0, rn), value=0.0)\n",
+    "            return cos_padded, sin_padded\n",
+    "\n",
+    "        pos_embeds = pad_rope(*pos_embeds)\n",
+    "        pos_embeds_for_freq = pad_rope(*pos_embeds_for_freq)\n",
+    "\n",
+    "        for time_transformer, freq_transformer in self.layers:\n",
+    "            hidden_states = time_transformer(\n",
+    "                hidden_states,\n",
+    "                position_embeddings=pos_embeds,\n",
+    "                attention_mask=None,\n",
+    "            )\n",
+    "            hidden_states = freq_transformer(\n",
+    "                hidden_states,\n",
+    "                position_embeddings=pos_embeds_for_freq,\n",
+    "                attention_mask=None,\n",
+    "            )\n",
+    "\n",
+    "        hidden_states = hidden_states[:, :t, :n, :]\n",
+    "\n",
+    "        return self.final_norm(hidden_states)\n",
+    "\n",
+    "\n",
+    "class BSRoformerForMaskedEstimation(BSRoformerPreTrainedModel):\n",
+    "    def __init__(self, config: BSRoformerConfig):\n",
+    "        super().__init__(config)\n",
+    "        self.config = config\n",
+    "        self.model = BSRoformerModel(config)\n",
+    "        self.mask_estimators = nn.ModuleList([MaskEstimator(config) for _ in range(config.num_stems)])\n",
+    "\n",
+    "        # STFT parameters\n",
+    "        self.stft_kwargs = dict(\n",
+    "            n_fft=config.stft_n_fft,\n",
+    "            hop_length=config.stft_hop_length,\n",
+    "            win_length=config.stft_win_length,\n",
+    "            normalized=False,\n",
+    "        )\n",
+    "        self.register_buffer(\"stft_window\", torch.hann_window(config.stft_win_length), persistent=False)\n",
+    "\n",
+    "        freqs = config.stft_n_fft // 2 + 1\n",
+    "        assert sum(config.freqs_per_bands) == freqs, f\"Sum of freqs_per_bands must be {freqs}\"\n",
+    "        self.wave_channels = config.num_input_channel\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        raw_audio: torch.Tensor,\n",
+    "        target: Optional[torch.Tensor] = None,\n",
+    "        return_loss_breakdown: bool = False,\n",
+    "    ):\n",
+    "        \"\"\"\n",
+    "        Args:\n",
+    "            raw_audio (`torch.Tensor` of shape `(batch, channels, time)`):\n",
+    "                The raw audio waveform.\n",
+    "            target (`torch.Tensor`, *optional*, shape `(batch, num_stems, channels, time)`):\n",
+    "                The target audio waveform for loss calculation.\n",
+    "            return_loss_breakdown (`bool`, *optional*, defaults to `False`):\n",
+    "                Whether to return the breakdown of the loss components.\n",
+    "\n",
+    "        Returns:\n",
+    "            torch.Tensor (`torch.Tensor` of shape `(batch, num_stems, channels, time)`):\n",
+    "            The reconstructed audio waveform.\n",
+    "        \"\"\"\n",
+    "        device = raw_audio.device\n",
+    "\n",
+    "        # 1. STFT: Convert audio to spectrogram\n",
+    "        with torch.autocast(device_type=device.type, enabled=False):\n",
+    "            b, c, t = raw_audio.shape  # batch, channel, time\n",
+    "            raw_audio_packed = rearrange(raw_audio, \"b c t -> (b c) t\")\n",
+    "            stft_repr = torch.stft(\n",
+    "                raw_audio_packed,\n",
+    "                **self.stft_kwargs,\n",
+    "                window=self.stft_window,\n",
+    "                return_complex=True,\n",
+    "            )\n",
+    "            stft_repr = torch.view_as_real(stft_repr)  # (b, c, t) -> (b, c, f, t, 2)\n",
+    "            stft_repr = rearrange(stft_repr, \"(b c) f t T -> b c f t T\", c=c)\n",
+    "            # Merge frequency, channel, and complex dimensions for the model\n",
+    "            stft_repr_merged = rearrange(stft_repr, \"b c f t T -> b t (f c T)\")\n",
+    "\n",
+    "        # 2. Model Processing\n",
+    "        hidden_states = self.model(stft_repr_merged)\n",
+    "\n",
+    "        # 3. Mask Estimation\n",
+    "        # (b, t, d) -> (b, n, t, (f c 2)) where n is num_stems\n",
+    "        mask = torch.stack([fn(hidden_states) for fn in self.mask_estimators], dim=1)\n",
+    "        mask = rearrange(mask, \"b n t (f c T) -> b n c f t T\", T=2, c=c)\n",
+    "        mask = mask.to(dtype=torch.float32)\n",
+    "\n",
+    "        # 4. Mask Application\n",
+    "        with torch.autocast(device_type=device.type, enabled=False):\n",
+    "            stft_repr_expanded = rearrange(stft_repr, \"b c f t T -> b 1 c f t T\")\n",
+    "            stft_repr_complex = torch.view_as_complex(stft_repr_expanded)\n",
+    "            mask_complex = torch.view_as_complex(mask)\n",
+    "            masked_stft = stft_repr_complex * mask_complex\n",
+    "\n",
+    "            # 5. iSTFT: Convert masked spectrogram back to audio\n",
+    "            # (b, n, c, f, t) -> ((b n c), f, t)\n",
+    "            masked_stft = rearrange(masked_stft, \"b n c f t -> (b n c) f t\")\n",
+    "            recon_audio = torch.istft(\n",
+    "                masked_stft,\n",
+    "                **self.stft_kwargs,\n",
+    "                window=self.stft_window,\n",
+    "                return_complex=False,\n",
+    "                length=raw_audio.shape[-1],\n",
+    "            )\n",
+    "            # ((b n c), t) -> (b, n, c, t)\n",
+    "            recon_audio = rearrange(recon_audio, \"(b n c) t -> b n c t\", c=self.wave_channels, n=self.config.num_stems)\n",
+    "\n",
+    "        if target is None:\n",
+    "            return recon_audio\n",
+    "\n",
+    "        # 6. Loss Calculation\n",
+    "        # Ensure target has the same length as the reconstructed audio\n",
+    "        target = target[..., : recon_audio.shape[-1]]\n",
+    "\n",
+    "        loss = F.l1_loss(recon_audio, target)\n",
+    "\n",
+    "        return loss\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f0f2b263",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model_config = BSRoformerConfig(\n",
+    "#     hidden_size=64,\n",
+    "#     depth=1,\n",
+    "#     num_input_channel=2,\n",
+    "#     num_stems=4,\n",
+    "#     intermediate_size=64 * 2,\n",
+    "#     time_transformer_depth=1,\n",
+    "#     freq_transformer_depth=1,\n",
+    "#     num_attention_heads=8,\n",
+    "#     num_key_value_heads=2,\n",
+    "#     #\n",
+    "#     mask_estimator_depth=1,\n",
+    "# )\n",
+    "# model = BSRoformerForMaskedEstimation(model_config)\n",
+    "\n",
+    "# dummy_input = torch.randn(6, 2, 44100 * 6)\n",
+    "# output = model(dummy_input)\n",
+    "\n",
+    "# dummy_targets = torch.randn(6, 4, 2, 44100 * 6)\n",
+    "# loss = model(dummy_input, target=dummy_targets)\n",
+    "\n",
+    "# dummy_song = torch.randn(2, 44100 * 30)\n",
+    "# result = inference_one_with_model(\n",
+    "#     model,\n",
+    "#     dummy_song,\n",
+    "#     chunk_size=44100 * 6,\n",
+    "#     overlap_size=44100 * 3,\n",
+    "#     gap_size=44100 * 1,\n",
+    "# )\n",
+    "\n",
+    "# del model, model_config, dummy_input, output, dummy_targets, loss"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9d26ff61",
+   "metadata": {},
+   "source": [
+    "## 实例化 Datasets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f4d791c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataset = AugmentDataset(\n",
+    "    data_path=[\n",
+    "        \"/mnt/sda/data/20250826_MUSDB18HQ/train\",\n",
+    "        \"/mnt/sda/data/20250826_MUSDB18HQ/test\",\n",
+    "        # \"/mnt/sda/data/20250902_DSD100/datas\",\n",
+    "    ],\n",
+    "    wave_chunk_size=44100 * 6,\n",
+    "    stem_names=[\"bass\", \"drums\", \"other\", \"vocals\"],\n",
+    ")\n",
+    "val_dataset = ValidationDataset(\n",
+    "    data_path=\"/mnt/sda/data/20250826_MUSDB18HQ/valid\",\n",
+    "    stem_names=[\"bass\", \"drums\", \"other\", \"vocals\"],\n",
+    ")\n",
+    "\n",
+    "train_loader = DataLoader(\n",
+    "    train_dataset,\n",
+    "    batch_size=18,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    "    persistent_workers=True if num_workers > 0 else False,\n",
+    "    prefetch_factor=4 if num_workers > 0 else None,\n",
+    ")\n",
+    "val_loader = DataLoader(\n",
+    "    val_dataset,\n",
+    "    batch_size=1,\n",
+    "    num_workers=num_workers,\n",
+    "    pin_memory=True,\n",
+    "    persistent_workers=True if num_workers > 0 else False,\n",
+    "    shuffle=False,\n",
+    "    prefetch_factor=4 if num_workers > 0 else None,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "21701211",
+   "metadata": {},
+   "source": [
+    "## Lightning"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23cda886",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def compute_sdr(target, estimate):\n",
+    "    target_np = target.float().cpu().numpy()\n",
+    "    estimate_np = estimate.float().cpu().numpy()\n",
+    "\n",
+    "    sdr_list = []\n",
+    "\n",
+    "    for this_target, this_estimate in zip(target_np, estimate_np):\n",
+    "        channel_sdrs = []\n",
+    "        for this_channel_target, this_channel_estimate in zip(this_target, this_estimate):\n",
+    "            signal_power = np.sum(this_channel_target ** 2)\n",
+    "            noise_power = np.sum((this_channel_target - this_channel_estimate) ** 2)\n",
+    "\n",
+    "            if noise_power == 0:\n",
+    "                sdr = float('inf')\n",
+    "            else:\n",
+    "                sdr = 10 * np.log10(signal_power / noise_power)\n",
+    "\n",
+    "            # sdr_list.append(sdr)\n",
+    "            channel_sdrs.append(sdr)\n",
+    "\n",
+    "        channel_sdr_mean = np.mean(channel_sdrs)\n",
+    "        sdr_list.append(channel_sdr_mean)\n",
+    "\n",
+    "    return sdr_list\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2e5002b1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class LightningModel(BaseModule):\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        model,\n",
+    "        lr_config: LearningRateConfig,\n",
+    "        training_config: TrainingConfig,\n",
+    "    ):\n",
+    "        super().__init__(\n",
+    "            model,\n",
+    "            lr_config,\n",
+    "            training_config,\n",
+    "        )\n",
+    "\n",
+    "        self.validation_sdr_results = []\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.model(x)\n",
+    "\n",
+    "    def training_step(self, batch, batch_idx):\n",
+    "        target_stems, mixed_audio = batch\n",
+    "        # target_stems: (batch, stems, channels, time)\n",
+    "        # mixed_audio: (batch, channels, time)\n",
+    "\n",
+    "        loss = self.model(mixed_audio, target=target_stems)\n",
+    "\n",
+    "        grad_norm = torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=4.0)\n",
+    "\n",
+    "        self.log('train/grad_norm', grad_norm.item(), on_step=True, on_epoch=False, sync_dist=True)\n",
+    "        self.log('train/loss', loss, on_step=True, on_epoch=False, sync_dist=True)\n",
+    "\n",
+    "        return loss\n",
+    "\n",
+    "    def validation_step(self, batch, batch_idx):\n",
+    "        target_stems, mixed_audio = batch\n",
+    "\n",
+    "        batch_size = mixed_audio.shape[0]\n",
+    "        batch_sdr_scores = []\n",
+    "\n",
+    "        for i in range(batch_size):\n",
+    "            single_mixed = mixed_audio[i]  # (channels, time)\n",
+    "            single_target = target_stems[i]  # (stems, channels, time)\n",
+    "\n",
+    "            with torch.no_grad():\n",
+    "                predicted_stems = inference_one_with_model(\n",
+    "                    self.model,\n",
+    "                    single_mixed,\n",
+    "                    chunk_size=44100 * 6,\n",
+    "                    overlap_size=44100 * 3,\n",
+    "                    gap_size=0,\n",
+    "                )  # (stems, channels, time)\n",
+    "\n",
+    "            sdr = compute_sdr(single_target, predicted_stems)\n",
+    "            batch_sdr_scores.append(sdr)\n",
+    "\n",
+    "        sdrs = np.array(batch_sdr_scores)\n",
+    "        sdrs = sdrs.mean(axis=0)\n",
+    "\n",
+    "        self.validation_sdr_results.append(sdrs)\n",
+    "\n",
+    "        return {\n",
+    "            \"val/sdr\": sdrs,\n",
+    "        }\n",
+    "\n",
+    "    def on_validation_epoch_end(self):\n",
+    "        if len(self.validation_sdr_results) > 0:\n",
+    "            avg_sdrs = np.mean(self.validation_sdr_results, axis=0)\n",
+    "            self.log('val/sdr', avg_sdrs.mean(), on_step=False, on_epoch=True, sync_dist=True, prog_bar=True)\n",
+    "            for i, one in enumerate(avg_sdrs):\n",
+    "                self.log(f'val/sdr_stem_{i}', one, on_step=False, on_epoch=True, sync_dist=True)\n",
+    "\n",
+    "            self.validation_sdr_results.clear()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e13d2d53",
+   "metadata": {},
+   "source": [
+    "## 配置与实例化"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c31b1d32",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pl_utils import LearningRateConfig, TrainingConfig\n",
+    "\n",
+    "\n",
+    "learning_rate_config = LearningRateConfig(\n",
+    "    lr_warmup_steps=400,\n",
+    "    lr_initial=1e-5,\n",
+    "    lr_max=5e-4,\n",
+    "    lr_end=5e-4,\n",
+    "    max_steps=20000,\n",
+    ")\n",
+    "\n",
+    "training_config = TrainingConfig(\n",
+    "    optimizer='adamw',\n",
+    "    optimizer_args={\n",
+    "        'betas': (0.9, 0.95),\n",
+    "        'weight_decay': 1e-2,\n",
+    "        \"fused\": True,\n",
+    "    },\n",
+    "    excluded_from_weight_decay=[\"bias\", \"norm\", \"embed\", \"scale\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13030935",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model_config = BSRoformerConfig(\n",
+    "    hidden_size=256,\n",
+    "    depth=3,\n",
+    "    num_input_channel=2,\n",
+    "    num_stems=4,\n",
+    "    intermediate_size=256 * 2,\n",
+    "    time_transformer_depth=1,\n",
+    "    freq_transformer_depth=1,\n",
+    "    num_attention_heads=8,\n",
+    "    num_key_value_heads=4,\n",
+    "    #\n",
+    "    mask_estimator_depth=1,\n",
+    "    multi_stft_loss_weight=0.0,\n",
+    ")\n",
+    "model = BSRoformerForMaskedEstimation(model_config)\n",
+    "\n",
+    "pl_model = LightningModel(\n",
+    "    model,\n",
+    "    lr_config=learning_rate_config,\n",
+    "    training_config=training_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9a430e4f",
+   "metadata": {},
+   "source": [
+    "## 正式训练"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d52e16e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lightning.pytorch.utilities.model_summary import summarize\n",
+    "\n",
+    "summarize(pl_model, max_depth=2)\n",
+    "\n",
+    "model.model.compile(options={\"shape_padding\": True})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6ff9af11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import lightning.pytorch as L\n",
+    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
+    "from lightning.pytorch.loggers import TensorBoardLogger\n",
+    "from pl_utils.lightning import format_next_version_name\n",
+    "from lightning.pytorch.strategies import DDPStrategy\n",
+    "\n",
+    "name = \"准备收尾。3层小模型，batch18\"\n",
+    "logger = TensorBoardLogger(save_dir=\"./\", version=format_next_version_name(name))\n",
+    "\n",
+    "checkpoint_callback = ModelCheckpoint(\n",
+    "    auto_insert_metric_name=True,\n",
+    "    save_top_k=1,\n",
+    "    monitor=\"val/sdr\",\n",
+    "    mode=\"max\",\n",
+    "    every_n_epochs=1,\n",
+    "    save_weights_only=True,\n",
+    "    # save_last=\"link\",\n",
+    "    save_on_train_epoch_end=False,\n",
+    "    save_last=True,\n",
+    ")\n",
+    "\n",
+    "trainer = L.Trainer(\n",
+    "    logger=logger,\n",
+    "    accelerator='gpu',\n",
+    "    # max_epochs=16,\n",
+    "    strategy=DDPStrategy(find_unused_parameters=False),\n",
+    "    precision='16-mixed',\n",
+    "    # accumulate_grad_batches=4,\n",
+    "    max_steps=200000,\n",
+    "    val_check_interval=500,\n",
+    "    log_every_n_steps=4,\n",
+    "    default_root_dir=\"./\",\n",
+    "    #\n",
+    "    callbacks=[checkpoint_callback],\n",
+    "    # enable_checkpointing=False,\n",
+    "    #\n",
+    "    num_sanity_val_steps=0,\n",
+    "    # fast_dev_run=True,\n",
+    "    # enable_checkpointing=False,\n",
+    "    enable_model_summary=True,\n",
+    ")\n",
+    "\n",
+    "trainer.fit(pl_model, train_loader, val_loader)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f9304c3e",
+   "metadata": {},
+   "source": [
+    "## 提前退出"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5e11d871",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "from IPython import get_ipython\n",
+    "\n",
+    "\n",
+    "# 如果是脚本而不是jupyter notebook，此时就该退出了\n",
+    "try:\n",
+    "    shell = get_ipython()\n",
+    "    if shell is None:\n",
+    "        sys.exit()\n",
+    "except:\n",
+    "    sys.exit()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8bca044c",
+   "metadata": {},
+   "source": [
+    "## 加载与推理"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2f65245d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pl_model = LightningModel.load_from_checkpoint(\n",
+    "    \"lightning_logs/version_029_可学习残差（策略为共享一个参数）/checkpoints/last.ckpt\",\n",
+    "    model=model,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f865e1f5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "waves, mixed_wave = val_dataset[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "badd73dd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with torch.inference_mode():\n",
+    "    predicted_stems = inference_one_with_model(\n",
+    "        pl_model.model,\n",
+    "        torch.tensor(mixed_wave).to(\"cuda\"),\n",
+    "        chunk_size=44100 * 6,\n",
+    "        overlap_size=44100 * 3,\n",
+    "    )  # (stems, channels, time)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a1fa8bbd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "predicted_stems.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "751c4974",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "os.makedirs(\"./outputs\", exist_ok=True)\n",
+    "\n",
+    "for i in range(predicted_stems.shape[0]):\n",
+    "    import soundfile as sf\n",
+    "\n",
+    "    sf.write(f\"./outputs/predicted_stem_{i}.wav\", predicted_stems[i].cpu().numpy().T, 44100)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d934415c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sf.write(\"./outputs/mixed.wav\", mixed_wave.T, 44100)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "20250820_bs-roformer",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.13.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}