first commit

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.csv filter=lfs diff=lfs merge=lfs -text
+*.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+
+checkpoints
+lightning_logs
+music_audioset_epoch_15_esc_90.14.pt

README.md

@@ -0,0 +1 @@
+# ClearSep

data_engine_processing/calc_resdr.py

@@ -0,0 +1,104 @@
+import os
+import csv
+import glob
+from tqdm import tqdm
+import torch
+import torchaudio
+from torchmetrics.audio import ScaleInvariantSignalDistortionRatio, SignalDistortionRatio
+
+
+def calculate_sdr_and_sisdr(original_audio_path, separated_audio_paths):
+    """
+    计算叠加的音频与原始音频之间的 SDR 和 SI-SDR。
+
+    参数:
+    - original_audio_path: str, 原始音频文件路径。
+    - separated_audio_paths: List[str], 分割后的音频片段文件路径列表。
+
+    返回:
+    - sdr: float, SDR 值。
+    - sisdr: float, SI-SDR 值。
+    """
+    # 加载原始音频
+    original_waveform, sample_rate = torchaudio.load(original_audio_path)
+
+    # 初始化叠加的音频波形
+    combined_waveform = None
+
+    # 加载并叠加分割的音频片段
+    for path in separated_audio_paths:
+        separated_waveform, _ = torchaudio.load(path)
+
+        # 对齐片段长度
+        min_length = min(original_waveform.size(1), separated_waveform.size(1))
+        separated_waveform = separated_waveform[:, :min_length]
+
+        # 初始化或叠加音频
+        if combined_waveform is None:
+            combined_waveform = separated_waveform
+        else:
+            combined_waveform = combined_waveform[:, :min_length] + separated_waveform
+
+    # 确保合并后的音频和原始音频的长度一致
+    min_length = min(original_waveform.size(1), combined_waveform.size(1))
+    original_waveform = original_waveform[:, :min_length]
+    combined_waveform = combined_waveform[:, :min_length]
+
+    # 计算 SI-SDR
+    sisdr_metric = ScaleInvariantSignalDistortionRatio()
+    sisdr = sisdr_metric(combined_waveform, original_waveform).item()
+
+    # 计算 SDR
+    sdr_metric = SignalDistortionRatio()
+    sdr = sdr_metric(combined_waveform, original_waveform).item()
+
+    # print(f"SI-SDR between original and combined audio: {sisdr} dB")
+    # print(f"SDR between original and combined audio: {sdr} dB")
+
+    return sdr, sisdr
+
+
+if __name__ == "__main__":
+    # 示例: 指定原始音频和分割后的音频片段路径
+    # original_audio_path = "path_to_original_audio.wav"
+    # separated_audio_paths = [
+    #     "path_to_segment_1.wav",
+    #     "path_to_segment_2.wav",
+    #     "path_to_segment_3.wav",
+    # ]
+
+    # # 计算 SDR 和 SI-SDR
+    # sdr, sisdr = calculate_sdr_and_sisdr(original_audio_path, separated_audio_paths)
+    
+    
+    dset = 'balanced_train_segments'
+    # dset = 'eval_segments'
+    
+    src_data_root = r'/data/sound/audioset/audios_32k'
+    sep_data_root = r'data_engine_infer/audioset_separation_child_label'
+    
+    writer = csv.writer(open(os.path.join(sep_data_root, dset + '.csv'), 'w'))
+    writer.writerow(['video', 'sdr', 'sisdr'])
+    for video_path in tqdm(glob.glob(os.path.join(sep_data_root, dset, '*'))):
+        video = video_path.split('/')[-1]
+        original_audio_path = os.path.join(src_data_root, dset, video + '.wav')
+        separated_audio_paths = glob.glob(video_path + '/*')
+        sdr, sisdr = calculate_sdr_and_sisdr(original_audio_path, separated_audio_paths)
+        writer.writerow([video, f'{sdr:.3f}', f'{sisdr:.3f}'])
+
+
+    # dset = 'unbalanced_train_segments'
+    
+    # src_data_root = r'/data/sound/audioset/audios_32k'
+    # sep_data_root = r'data_engine_infer/audioset_separation_child_label'
+    
+    # writer = csv.writer(open(os.path.join(sep_data_root, dset + '.csv'), 'w'))
+    # writer.writerow(['video', 'sdr', 'sisdr'])
+    # for video_path in tqdm(glob.glob(os.path.join(sep_data_root, dset, '*', '*'))):
+    #     part = video_path.split('/')[-2]
+    #     video = video_path.split('/')[-1]
+    #     original_audio_path = os.path.join(src_data_root, dset, part, video + '.wav')
+    #     separated_audio_paths = glob.glob(video_path + '/*')
+    #     sdr, sisdr = calculate_sdr_and_sisdr(original_audio_path, separated_audio_paths)
+    #     writer.writerow([video, f'{sdr:.3f}', f'{sisdr:.3f}'])
+

data_engine_processing/filter_threshold.py

@@ -0,0 +1,35 @@
+import os
+import csv
+import glob
+from tqdm import tqdm
+
+
+audioset_labels = set()
+reader = csv.reader(open('metadata/audioset/class_labels_indices.csv', 'r'))
+next(reader)
+for item in reader:
+    assert len(item) == 3
+    audioset_labels.add(item[-1])
+    
+sdr_dict = dict()
+reader = csv.reader(open('data_engine_infer/audioset_separation_child_label/balanced_train_segments.csv', 'r'))
+next(reader)
+for item in reader:
+    sdr_dict[item[0]] = min(float(item[1]), float(item[2]))
+
+
+writer_10 = csv.writer(open('data_engine_infer/audioset_separation_child_label/audioset_bal_data_engine_th10.csv', 'w'))
+writer_15 = csv.writer(open('data_engine_infer/audioset_separation_child_label/audioset_bal_data_engine_th15.csv', 'w'))
+
+for video_path in tqdm(glob.glob('data_engine_infer/audioset_separation_child_label/balanced_train_segments/*')):
+    video = video_path.split('/')[-1]
+    assert video in sdr_dict
+    if sdr_dict[video] < 10.0:
+        continue
+    for wav_path in glob.glob(video_path + '/*.wav'):
+        label = wav_path.split('/')[-1][:-4]
+        assert label in audioset_labels, f"label: {label}, wav_path: {wav_path}"
+        writer_10.writerow([wav_path, label])
+        if sdr_dict[video] >= 15.0:
+            writer_15.writerow([wav_path, label])
+

dataset.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main
+@File    ：dataset_online.py
+@IDE     ：PyCharm
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2023/11/1 下午6:47
+'''
+import os
+import random
+
+import torch
+import torchaudio
+import torchaudio.transforms as AT
+import csv
+import json
+import numpy as np
+import librosa
+
+
+def labels2caption(labels):
+    prefix = "The sound of " if len(labels) == 1 else "The sounds of "
+    caption = prefix + ', '.join(labels)
+    return caption
+
+
+class CLAPSepDataSet(torch.utils.data.Dataset):  # type: ignore
+
+    def __init__(self, data_list, dset='', silence_rate=0.05, chunk_dur=10, sr=None, resample_rate=None):
+        assert dset in ['train', 'val'], \
+            "`dset` must be one of ['train', 'val']"
+        self.dset = dset
+        self.silence_rate = silence_rate
+        self.chunk_dur = chunk_dur
+        self.data_meta = dict()
+        self.text_dict = dict()
+        with open(data_list, 'r', encoding='utf-8') as d:
+            reader = csv.reader(d, skipinitialspace=True)
+            for row in reader:
+                assert os.path.exists(row[0])
+                self.data_meta[row[0]] = row[1:]
+                label = ', '.join(row[1:])
+                if label not in self.text_dict:
+                    self.text_dict[label] = []
+                self.text_dict[label].append(row[0])
+        # self.data_meta.pop('file_name')
+        self.augmentation = torchaudio.transforms.SpeedPerturbation(48000, [0.9, 1.1])
+
+        self.data_names = list(self.data_meta.keys())
+        if dset == 'val':
+            self.noise_names = []
+            for name in self.data_names:
+                noise_name = self.choose_other_samples(', '.join(self.data_meta[name]), 1)[0]
+                self.noise_names.append(noise_name)
+
+        if resample_rate is not None:
+            self.resampler = AT.Resample(sr, resample_rate)
+            self.sr = sr
+            self.resample_rate = resample_rate
+        else:
+            self.sr = sr
+
+    def __len__(self):
+        return len(self.data_names)
+
+    def choose_other_samples(self, target_text, num):
+        candidates = list(self.text_dict.keys())
+        candidates.remove(target_text)
+        chosen_text = random.sample(candidates, num)
+        chosen_samples = [random.choice(self.text_dict[text]) for text in chosen_text]
+        return chosen_samples
+
+    def load_wav(self, path):
+        max_length = self.sr * self.chunk_dur
+        wav = librosa.core.load(path, sr=self.sr)[0]
+        if len(wav) > max_length:
+            wav = wav[0:max_length]
+
+        # pad audio to max length, 10s for AudioCaps
+        if len(wav) < max_length:
+            wav = np.pad(wav, (0, max_length - len(wav)), 'constant')
+        return wav
+
+    def __getitem__(self, idx):
+        tgt_name = self.data_names[idx]
+        if self.dset =='train':
+            noise_name = tgt_name
+            while set(self.data_meta[noise_name]) & set(self.data_meta[tgt_name]):
+                noise_name = random.choice(self.data_names)
+        else:
+            noise_name = self.noise_names[idx]
+
+        snr = torch.zeros((1,))
+        # snr = (torch.rand((1,)) * 10 - 5) if self.dset == 'train' else torch.zeros((1,))
+        tgt = torch.tensor(self.load_wav(tgt_name)).unsqueeze(0)
+        noise = torch.tensor(self.load_wav(noise_name)).unsqueeze(0)
+        # assert not torch.isnan(tgt).any()
+        # assert not torch.isnan(noise).any()
+        mixed = torchaudio.functional.add_noise(tgt, noise, snr=snr)
+        assert not torch.isnan(mixed).any(), f"tgt: {tgt_name}, noise: {noise_name}"
+        pos_sample, _ = self.augmentation(self.resampler(tgt.squeeze()))
+        neg_sample, _ = self.augmentation(self.resampler(noise.squeeze()))
+
+        max_value = torch.max(torch.abs(mixed))
+        if max_value > 1:
+            tgt *= 0.9 / max_value
+            mixed *= 0.9 / max_value
+
+        tgt = tgt.squeeze()
+        mixed = mixed.squeeze()
+        tgt_cap = labels2caption(self.data_meta[tgt_name])
+        neg_cap = labels2caption(self.data_meta[noise_name])
+        mixed_resample = self.resampler(mixed)
+        
+        # silence query
+        if self.dset =='train' and random.random() < self.silence_rate:
+            other_name = tgt_name
+            while set(self.data_meta[other_name]) & (set(self.data_meta[tgt_name]) | set(self.data_meta[noise_name])):
+                other_name = random.choice(self.data_names)
+            tgt = torch.zeros_like(mixed)
+            neg_cap = labels2caption(self.data_meta[tgt_name] + self.data_meta[noise_name])
+            tgt_cap = labels2caption(self.data_meta[other_name])
+            pos_sample, _ = self.augmentation(self.resampler(torch.tensor(self.load_wav(other_name))))
+            neg_sample, _ = self.augmentation(mixed_resample)
+
+        return mixed, mixed_resample, tgt_cap, neg_cap, tgt, self.pad_or_trim(pos_sample), self.pad_or_trim(neg_sample)
+
+    def pad_or_trim(self, wav_in):
+        target_len = 48000 * self.chunk_dur
+        if wav_in.size(0) < target_len:
+            wav_in = torch.nn.functional.pad(wav_in, (0, target_len - wav_in.size(0)))
+        elif wav_in.size(0) > target_len:
+            wav_in = wav_in[:target_len]
+        max_value = torch.max(torch.abs(wav_in))
+        if max_value > 1:
+            wav_in *= 0.9 / max_value
+        return wav_in
+
+
+class CLAPSepDataEngineDataSet(torch.utils.data.Dataset):  # type: ignore
+
+    def __init__(self, data_list, dset='', data_engine_json='', silence_rate=0.05, chunk_dur=10, sr=None, resample_rate=None):
+        assert dset in ['train', 'val'], \
+            "`dset` must be one of ['train', 'val']"
+        self.dset = dset
+        self.silence_rate = silence_rate
+        self.chunk_dur = chunk_dur
+        self.data_meta = dict()
+        with open(data_list, 'r', encoding='utf-8') as d:
+            reader = csv.reader(d, skipinitialspace=True)
+            for row in reader:
+                assert os.path.exists(row[0]), row[0]
+                self.data_meta[row[0]] = row[1:]
+        # self.data_meta.pop('file_name')
+        self.augmentation = torchaudio.transforms.SpeedPerturbation(48000, [0.9, 1.1])
+
+        self.data_names = list(self.data_meta.keys())
+        if dset == 'val':
+            self.noise_names = []
+            for name in self.data_names:
+                noise_name = name
+                while set(self.data_meta[noise_name]) & set(self.data_meta[name]):
+                    noise_name = random.choice(self.data_names)
+                self.noise_names.append(noise_name)
+        
+        self.data_engine_dict = {}
+        if os.path.exists(data_engine_json):
+            self.data_engine_dict = json.load(open(data_engine_json, 'r'))
+
+        if resample_rate is not None:
+            self.resampler = AT.Resample(sr, resample_rate)
+            self.sr = sr
+            self.resample_rate = resample_rate
+        else:
+            self.sr = sr
+
+    def __len__(self):
+        return len(self.data_names)
+
+    def load_wav(self, path):
+        max_length = self.sr * self.chunk_dur
+        wav = librosa.core.load(path, sr=self.sr)[0]
+        if len(wav) > max_length:
+            wav = wav[0:max_length]
+
+        # pad audio to max length, 10s for AudioCaps
+        if len(wav) < max_length:
+            wav = np.pad(wav, (0, max_length - len(wav)), 'constant')
+        return wav
+
+    def __getitem__(self, idx):
+        tgt_name = self.data_names[idx]
+        if self.dset =='train':
+            noise_name = tgt_name
+            while set(self.data_meta[noise_name]) & set(self.data_meta[tgt_name]):
+                noise_name = random.choice(self.data_names)
+        else:
+            noise_name = self.noise_names[idx]
+        
+        snr = torch.zeros((1,))
+        # snr = (torch.rand((1,)) * 10 - 5) if self.dset == 'train' else torch.zeros((1,))
+        tgt = torch.tensor(self.load_wav(tgt_name)).unsqueeze(0)
+        noise = torch.tensor(self.load_wav(noise_name)).unsqueeze(0)
+        # assert not torch.isnan(tgt).any()
+        # assert not torch.isnan(noise).any()
+        mixed = torchaudio.functional.add_noise(tgt, noise, snr=snr)
+        # assert not torch.isnan(mixed).any(), f"tgt: {tgt_name}, noise: {noise_name}"
+        
+        pos_sample, _ = self.augmentation(self.resampler(tgt.squeeze()))
+        noise = noise.squeeze()
+        
+        max_value = torch.max(torch.abs(mixed))
+        if max_value > 1:
+            tgt *= 0.9 / max_value
+            mixed *= 0.9 / max_value
+        
+        tgt = tgt.squeeze()
+        mixed = mixed.squeeze()
+        tgt_cap = labels2caption(self.data_meta[tgt_name])
+        neg_cap = labels2caption(self.data_meta[noise_name])
+        mixed_resample = self.resampler(mixed)
+        
+        # A(A1, A2) + B, A1 as target, A2 + B as noise
+        # video = tgt_name.split('/')[-1][:-4]
+        # if self.dset =='train' and video in self.data_engine_dict and random.random() > 0.5:
+        #     items = self.data_engine_dict[video]
+        #     tgt_idx = random.choice(range(0, len(items)))
+        #     tgt_item = items[tgt_idx]
+        #     items.pop(tgt_idx)
+        #     tgt = torch.tensor(self.load_wav(tgt_item[0]))
+        #     max_value = torch.max(torch.abs(tgt))
+        #     if max_value > 1:
+        #         tgt *= 0.9 / max_value
+        #     tgt_cap = tgt_item[1]
+        #     if len(items) > 0:
+        #         noises = [torch.tensor(self.load_wav(x[0])) for x in items]
+        #         noises.append(noise)
+        #         noise_caps = [neg_cap.replace('sound', 'sounds')] + [x[1] for x in items]
+        #         noise = torch.mean(torch.stack(noises, dim=0), dim=0)
+        #         neg_cap = ', '.join(noise_caps)
+        
+        # A(A1, A2), A1 as target, others as noise
+        video = tgt_name.split('/')[-1][:-4]
+        if self.dset =='train' and video in self.data_engine_dict and random.random() > 0.5:
+            mixed = tgt
+            mixed_resample = self.resampler(mixed)
+            items = self.data_engine_dict[video]
+            tgt_idx = random.choice(range(0, len(items)))
+            tgt_item = items[tgt_idx]
+            items.pop(tgt_idx)
+            tgt = torch.tensor(self.load_wav(tgt_item[0]))
+            max_value = torch.max(torch.abs(tgt))
+            if max_value > 1:
+                tgt *= 0.9 / max_value
+            tgt_cap = tgt_item[1]
+            if len(items) > 0:
+                noises = [torch.tensor(self.load_wav(x[0])) for x in items]
+                noise_caps = [x[1] for x in items]
+                noise = torch.mean(torch.stack(noises, dim=0), dim=0)
+                neg_cap = labels2caption(noise_caps)
+        
+        # silence query
+        elif self.dset =='train' and random.random() < self.silence_rate:
+            other_name = tgt_name
+            while set(self.data_meta[other_name]) & (set(self.data_meta[tgt_name]) | set(self.data_meta[noise_name])):
+                other_name = random.choice(self.data_names)
+            tgt = torch.zeros_like(mixed)
+            neg_cap = labels2caption(self.data_meta[tgt_name] + self.data_meta[noise_name])
+            tgt_cap = labels2caption(self.data_meta[other_name])
+            pos_sample, _ = self.augmentation(self.resampler(torch.tensor(self.load_wav(other_name))))
+            noise = mixed
+        
+        neg_sample, _ = self.augmentation(self.resampler(noise))
+
+        return mixed, mixed_resample, tgt_cap, neg_cap, tgt, self.pad_or_trim(pos_sample), self.pad_or_trim(neg_sample)
+
+    def pad_or_trim(self, wav_in):
+        target_len = 48000 * self.chunk_dur
+        if wav_in.size(0) < target_len:
+            wav_in = torch.nn.functional.pad(wav_in, (0, target_len - wav_in.size(0)))
+        elif wav_in.size(0) > target_len:
+            wav_in = wav_in[:target_len]
+        max_value = torch.max(torch.abs(wav_in))
+        if max_value > 1:
+            wav_in *= 0.9 / max_value
+        return wav_in
+

eval.py

@@ -0,0 +1,161 @@
+import os
+import random
+
+import torch
+import torchaudio
+import torchaudio.transforms as AT
+import csv
+import numpy as np
+import librosa
+import pandas as pd
+import laion_clap
+import soundfile as sf
+from model.CLAPSep import LightningModule
+from model.CLAPSep_decoder import HTSAT_Decoder
+import argparse
+import pytorch_lightning as pl
+from helpers import utils as local_utils
+
+
+class AudioCapsTest(torch.utils.data.Dataset):  # type: ignore
+
+    def __init__(self, eval_csv, input_dir, sr=32000,
+                 resample_rate=48000):
+        self.data_path = input_dir
+
+        self.data_names = []
+        self.data_caps = []
+        self.noise_names = []
+        self.noise_caps = []
+        with open(eval_csv, 'r') as d:
+            reader = csv.reader(d, skipinitialspace=True)
+            next(reader)
+            for row in reader:
+                self.data_names.append(row[0])
+                self.data_caps.append(row[1])
+                self.noise_names.append(row[2])
+                self.noise_caps.append(row[3])
+
+        if resample_rate is not None:
+            self.resampler = AT.Resample(sr, resample_rate)
+            self.sr = sr
+            self.resample_rate = resample_rate
+        else:
+            self.sr = sr
+
+    def __len__(self):
+        return len(self.data_names)
+
+    def load_wav(self, path):
+        max_length = self.sr * 10
+        wav = librosa.core.load(path, sr=self.sr)[0]
+        if len(wav) > max_length:
+            wav = wav[0:max_length]
+
+        # pad audio to max length, 10s for AudioCaps
+        if len(wav) < max_length:
+            # audio = torch.nn.functional.pad(audio, (0, self.max_length - audio.size(1)), 'constant')
+            wav = np.pad(wav, (0, max_length - len(wav)), 'constant')
+        return wav
+
+    def __getitem__(self, idx):
+
+        tgt_name = self.data_names[idx]
+        noise_name = self.noise_names[idx]
+        tgt_cap = self.data_caps[idx]
+        neg_cap = self.noise_caps[idx]
+
+        assert noise_name != tgt_name
+        snr = torch.ones((1,)) * 0
+        tgt = torch.tensor(self.load_wav(os.path.join(self.data_path, tgt_name))).unsqueeze(0)
+        noise = torch.tensor(self.load_wav(os.path.join(self.data_path, noise_name))).unsqueeze(0)
+        mixed = torchaudio.functional.add_noise(tgt, noise, snr=snr)
+
+        max_value = torch.max(torch.abs(mixed))
+        if max_value > 1:
+            tgt *= 0.9 / max_value
+            mixed *= 0.9 / max_value
+        
+        tgt = tgt.squeeze()
+        mixed = mixed.squeeze()
+        
+        return mixed, self.resampler(mixed), tgt_cap, neg_cap, tgt
+
+
+
+def main(args):
+    torch.set_float32_matmul_precision('highest')
+    # Load dataset
+    
+    data_test = AudioCapsTest(eval_csv=args.eval_csv,
+                              input_dir=args.input_dir,
+                              sr=args.sample_rate,
+                              resample_rate=48000)
+
+    test_loader = torch.utils.data.DataLoader(data_test,
+                                             batch_size=1,
+                                             num_workers=1,
+                                             pin_memory=True,
+                                             shuffle=False)
+
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-base', device='cpu')
+    clap_model.load_ckpt(args.clap_path)
+    decoder = HTSAT_Decoder(**args.model)
+    lightning_module = LightningModule(clap_model, decoder, lr=args.optim['lr'],
+                                       use_lora=args.lora,
+                                       rank=args.lora_rank,
+                                       nfft=args.nfft)
+    distributed_backend = "ddp"
+    trainer = pl.Trainer(
+        default_root_dir=os.path.join(args.exp_dir, 'checkpoint'),
+        devices=args.gpu_ids if args.use_cuda else "auto",
+        accelerator="gpu" if args.use_cuda else "cpu",
+        benchmark=False,
+        gradient_clip_val=5.0,
+        precision='bf16-mixed',
+        limit_train_batches=1.0,
+        max_epochs=args.epochs,
+        strategy=distributed_backend,
+        logger=False
+    )
+
+    # weight = torch.load(args.ckpt_path, map_location="cpu")
+    # lightning_module.load_state_dict(weight, strict=False)
+
+    # trainer.test(model=lightning_module, dataloaders=test_loader)
+    
+    trainer.test(model=lightning_module, dataloaders=test_loader, ckpt_path=args.ckpt_path)
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Data Params
+    parser.add_argument('exp_dir', type=str,
+                        default='experiments',
+                        help="Path to save checkpoints and logs.")
+    
+    parser.add_argument('--sample_rate', type=int, default=32000)
+    parser.add_argument('--ckpt_path', type=str, default='')
+    parser.add_argument('--eval_csv', type=str, default='')
+    parser.add_argument('--input_dir', type=str, default='')
+
+    parser.add_argument('--use_cuda', dest='use_cuda', action='store_true',
+                        help="Whether to use cuda")
+    parser.add_argument('--gpu_ids', nargs='+', type=int, default=None,
+                        help="List of GPU ids used for training. "
+                             "Eg., --gpu_ids 2 4. All GPUs are used by default.")
+
+    args = parser.parse_args()
+
+    # Set the random seed for reproducible experiments
+    pl.seed_everything(114514)
+    # Set up checkpoints
+    if not os.path.exists(args.exp_dir):
+        os.makedirs(args.exp_dir)
+
+    # Load model and training params
+    params = local_utils.Params(os.path.join(args.exp_dir, 'config.json'))
+    for k, v in params.__dict__.items():
+        vars(args)[k] = v
+    main(args)

eval.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+ckpt_path="experiments/ClearSep_audioset_32k/checkpoints/epoch=100-step=868000-val_loss=10.33.ckpt"
+
+
+python eval.py \
+  ./experiments/ClearSep_audioset_32k \
+  --sample_rate 32000 \
+  --ckpt_path $ckpt_path \
+  --eval_csv metadata/evaluation/audiocaps_label_eval.csv \
+  --input_dir /home/wangtianhao/data/sound/audiocaps/dataset/audios_32k/test \
+  --use_cuda \
+  --gpu_ids 0
+

experiments/ClearSep_audioset_32k/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e375acdc54da7338b97f80035f272c6e019377363c60df0f107f9a00dd8a5267
+size 1104

helpers/utils.py

@@ -0,0 +1,205 @@
+"""A collection of useful helper functions"""
+
+import os
+import logging
+import json
+
+import torch
+from torch.profiler import profile, record_function, ProfilerActivity
+import pandas as pd
+from torchmetrics.functional import(
+    scale_invariant_signal_noise_ratio as si_snr,
+    signal_noise_ratio as snr,
+    signal_distortion_ratio as sdr,
+    scale_invariant_signal_distortion_ratio as si_sdr)
+import matplotlib.pyplot as plt
+
+class Params():
+    """Class that loads hyperparameters from a json file.
+    Example:
+    ```
+    params = Params(json_path)
+    print(params.learning_rate)
+    params.learning_rate = 0.5  # change the value of learning_rate in params
+    ```
+    """
+
+    def __init__(self, json_path):
+        with open(json_path) as f:
+            params = json.load(f)
+            self.__dict__.update(params)
+
+    def save(self, json_path):
+        with open(json_path, 'w') as f:
+            json.dump(self.__dict__, f, indent=4)
+
+    def update(self, json_path):
+        """Loads parameters from json file"""
+        with open(json_path) as f:
+            params = json.load(f)
+            self.__dict__.update(params)
+
+    @property
+    def dict(self):
+        """Gives dict-like access to Params instance by `params.dict['learning_rate']"""
+        return self.__dict__
+
+def save_graph(train_metrics, test_metrics, save_dir):
+    metrics = [snr, si_snr]
+    results = {'train_loss': train_metrics['loss'],
+               'test_loss' : test_metrics['loss']}
+
+    for m_fn in metrics:
+        results["train_"+m_fn.__name__] = train_metrics[m_fn.__name__]
+        results["test_"+m_fn.__name__] = test_metrics[m_fn.__name__]
+
+    results_pd = pd.DataFrame(results)
+
+    results_pd.to_csv(os.path.join(save_dir, 'results.csv'))
+
+    fig, temp_ax = plt.subplots(2, 3, figsize=(15,10))
+    axs=[]
+    for i in temp_ax:
+        for j in i:
+            axs.append(j)
+
+    x = range(len(train_metrics['loss']))
+    axs[0].plot(x, train_metrics['loss'], label='train')
+    axs[0].plot(x, test_metrics['loss'], label='test')
+    axs[0].set(ylabel='Loss')
+    axs[0].set(xlabel='Epoch')
+    axs[0].set_title('loss',fontweight='bold')
+    axs[0].legend()
+
+    for i in range(len(metrics)):
+        axs[i+1].plot(x, train_metrics[metrics[i].__name__], label='train')
+        axs[i+1].plot(x, test_metrics[metrics[i].__name__], label='test')
+        axs[i+1].set(xlabel='Epoch')
+        axs[i+1].set_title(metrics[i].__name__,fontweight='bold')
+        axs[i+1].legend()
+
+    plt.tight_layout()
+    plt.savefig(os.path.join(save_dir, 'results.png'))
+    plt.close(fig)
+
+def set_logger(log_path):
+    """Set the logger to log info in terminal and file `log_path`.
+    In general, it is useful to have a logger so that every output to the terminal is saved
+    in a permanent file. Here we save it to `model_dir/train.log`.
+    Example:
+    ```
+    logging.info("Starting training...")
+    ```
+    Args:
+        log_path: (string) where to log
+    """
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+    logger.handlers.clear()
+
+    # Logging to a file
+    file_handler = logging.FileHandler(log_path)
+    file_handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s: %(message)s'))
+    logger.addHandler(file_handler)
+
+    # Logging to console
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(logging.Formatter('%(message)s'))
+    logger.addHandler(stream_handler)
+
+def load_checkpoint(checkpoint, model, optim=None, lr_sched=None, data_parallel=False):
+    """Loads model parameters (state_dict) from file_path.
+
+    Args:
+        checkpoint: (string) filename which needs to be loaded
+        model: (torch.nn.Module) model for which the parameters are loaded
+        data_parallel: (bool) if the model is a data parallel model
+    """
+    if not os.path.exists(checkpoint):
+        raise("File doesn't exist {}".format(checkpoint))
+
+    state_dict = torch.load(checkpoint)
+
+    if data_parallel:
+        state_dict['model_state_dict'] = {
+            'module.' + k: state_dict['model_state_dict'][k]
+            for k in state_dict['model_state_dict'].keys()}
+    model.load_state_dict(state_dict['model_state_dict'])
+
+    if optim is not None:
+        optim.load_state_dict(state_dict['optim_state_dict'])
+
+    if lr_sched is not None:
+        lr_sched.load_state_dict(state_dict['lr_sched_state_dict'])
+
+    return state_dict['epoch'], state_dict['train_metrics'], \
+           state_dict['val_metrics']
+
+def save_checkpoint(checkpoint, epoch, model, optim=None, lr_sched=None,
+                    train_metrics=None, val_metrics=None, data_parallel=False):
+    """Saves model parameters (state_dict) to file_path.
+
+    Args:
+        checkpoint: (string) filename which needs to be loaded
+        model: (torch.nn.Module) model for which the parameters are loaded
+        data_parallel: (bool) if the model is a data parallel model
+    """
+    if os.path.exists(checkpoint):
+        raise("File already exists {}".format(checkpoint))
+
+    model_state_dict = model.state_dict()
+    if data_parallel:
+        model_state_dict = {
+            k.partition('module.')[2]:
+            model_state_dict[k] for k in model_state_dict.keys()}
+
+    optim_state_dict = None if not optim else optim.state_dict()
+    lr_sched_state_dict = None if not lr_sched else lr_sched.state_dict()
+
+    state_dict = {
+        'epoch': epoch,
+        'model_state_dict': model_state_dict,
+        'optim_state_dict': optim_state_dict,
+        'lr_sched_state_dict': lr_sched_state_dict,
+        'train_metrics': train_metrics,
+        'val_metrics': val_metrics
+    }
+
+    torch.save(state_dict, checkpoint)
+
+def model_size(model):
+    """
+    Returns size of the `model` in millions of parameters.
+    """
+    num_train_params = sum(
+        p.numel() for p in model.parameters() if p.requires_grad)
+    return num_train_params / 1e6
+
+def run_time(model, inputs, profiling=False):
+    """
+    Returns runtime of a model in ms.
+    """
+    # Warmup
+    for _ in range(100):
+        output = model(*inputs)
+
+    with profile(activities=[ProfilerActivity.CPU],
+                 record_shapes=True) as prof:
+        with record_function("model_inference"):
+            output = model(*inputs)
+
+    # Print profiling results
+    if profiling:
+        print(prof.key_averages().table(sort_by="self_cpu_time_total",
+                                        row_limit=20))
+
+    # Return runtime in ms
+    return prof.profiler.self_cpu_time_total / 1000
+
+def format_lr_info(optimizer):
+    lr_info = ""
+    for i, pg in enumerate(optimizer.param_groups):
+        lr_info += " {group %d: params=%.5fM lr=%.1E}" % (
+            i, sum([p.numel() for p in pg['params']]) / (1024 ** 2), pg['lr'])
+    return lr_info
+

infer_data_engine.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+ckpt_path="experiments/ClearSep_audioset_32k/checkpoints/epoch=100-step=868000-val_loss=10.33.ckpt"
+
+# output audio save path: ClearSep/model/CLAPSep_infer.py:153~154
+
+python infer_data_engine_json.py \
+  ./experiments/ClearSep_audioset_32k \
+  --sample_rate 32000 \
+  --ckpt_path $ckpt_path \
+  --audioset_json metadata/data_engine_meta/child_label/bal_train_segments_multi_label.json \
+  --video2path_map_csv metadata/audioset/train_lst.csv \
+  --use_cuda \
+  --gpu_ids 0
+

infer_data_engine_json.py

@@ -0,0 +1,143 @@
+import os
+import json
+import random
+
+import torch
+import torchaudio
+import torchaudio.transforms as AT
+import csv
+import numpy as np
+import librosa
+import pandas as pd
+import laion_clap
+from model.CLAPSep_infer import LightningModule
+from model.CLAPSep_decoder import HTSAT_Decoder
+import argparse
+import pytorch_lightning as pl
+from helpers import utils as local_utils
+
+
+class AudioCapsTest(torch.utils.data.Dataset):  # type: ignore
+
+    def __init__(self, audioset_json, video2path_map_csv, sr=32000, resample_rate=48000):
+        self.data_names = []
+        self.data_labels = []
+        video2path = {}
+        for item in csv.reader(open(video2path_map_csv, 'r')):
+            video2path[item[0]] = item[-1]
+        
+        video2labels = json.load(open(audioset_json, 'r'))
+        for video, labels in video2labels.items():
+            if video in video2path:
+                video_path = video2path[video]
+                self.data_names.append(video_path)
+                self.data_labels.append(labels)
+
+        if resample_rate is not None:
+            self.resampler = AT.Resample(sr, resample_rate)
+            self.sr = sr
+            self.resample_rate = resample_rate
+        else:
+            self.sr = sr
+
+    def __len__(self):
+        return len(self.data_names)
+
+    def load_wav(self, path):
+        max_length = self.sr * 10
+        wav = librosa.core.load(path, sr=self.sr)[0]
+        if len(wav) > max_length:
+            wav = wav[0:max_length]
+
+        # pad audio to max length, 10s for AudioCaps
+        if len(wav) < max_length:
+            # audio = torch.nn.functional.pad(audio, (0, self.max_length - audio.size(1)), 'constant')
+            wav = np.pad(wav, (0, max_length - len(wav)), 'constant')
+        return wav
+
+    def __getitem__(self, idx):
+        tgt_name = self.data_names[idx]
+        tgt_labels = self.data_labels[idx]
+
+        mixed = torch.tensor(self.load_wav(tgt_name))
+
+        return mixed, self.resampler(mixed), '|'.join(tgt_labels), tgt_name
+
+
+
+def main(args):
+    torch.set_float32_matmul_precision('highest')
+    # Load dataset
+    
+    data_test = AudioCapsTest(audioset_json=args.audioset_json,
+                              video2path_map_csv=args.video2path_map_csv,
+                              sr=args.sample_rate,
+                              resample_rate=48000)
+
+    test_loader = torch.utils.data.DataLoader(data_test,
+                                             batch_size=1,
+                                             num_workers=1,
+                                             pin_memory=True,
+                                             shuffle=False)
+
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-base', device='cpu')
+    clap_model.load_ckpt(args.clap_path)
+    decoder = HTSAT_Decoder(**args.model)
+    lightning_module = LightningModule(clap_model, decoder, lr=args.optim['lr'],
+                                       use_lora=args.lora,
+                                       rank=args.lora_rank,
+                                       nfft=args.nfft)
+    distributed_backend = "ddp"
+    trainer = pl.Trainer(
+        default_root_dir=os.path.join(args.exp_dir, 'checkpoint'),
+        devices=args.gpu_ids if args.use_cuda else "auto",
+        accelerator="gpu" if args.use_cuda else "cpu",
+        benchmark=False,
+        gradient_clip_val=5.0,
+        precision='bf16-mixed',
+        limit_train_batches=1.0,
+        max_epochs=args.epochs,
+        strategy=distributed_backend,
+        logger=False
+    )
+
+    weights = torch.load(args.ckpt_path, map_location='cpu')
+    lightning_module.load_state_dict(weights, strict=False)
+
+    trainer.test(model=lightning_module, dataloaders=test_loader)
+
+    # trainer.test(model=lightning_module, dataloaders=test_loader, ckpt_path=args.ckpt_path)
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Data Params
+    parser.add_argument('exp_dir', type=str,
+                        default='experiments',
+                        help="Path to save checkpoints and logs.")
+    
+    parser.add_argument('--sample_rate', type=int, default=16000)
+    parser.add_argument('--ckpt_path', type=str, default='')
+    parser.add_argument('--audioset_json', type=str, default='')
+    parser.add_argument('--video2path_map_csv', type=str, default='')
+
+    parser.add_argument('--use_cuda', dest='use_cuda', action='store_true',
+                        help="Whether to use cuda")
+    parser.add_argument('--gpu_ids', nargs='+', type=int, default=None,
+                        help="List of GPU ids used for training. "
+                             "Eg., --gpu_ids 2 4. All GPUs are used by default.")
+
+    args = parser.parse_args()
+
+    # Set the random seed for reproducible experiments
+    pl.seed_everything(114514)
+    # Set up checkpoints
+    if not os.path.exists(args.exp_dir):
+        os.makedirs(args.exp_dir)
+
+    # Load model and training params
+    params = local_utils.Params(os.path.join(args.exp_dir, 'config.json'))
+    for k, v in params.__dict__.items():
+        vars(args)[k] = v
+    main(args)

metadata/audioset/audioset_label.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e340235d4af6b34442d948bc41e2b2fadf90d7e4109702fb8520c4cc6c4b8f62
+size 69661293

metadata/audioset/class_labels_indices.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cdd1049833c4b86127c2773ac0d14a2754b6a6d0d1798002ed5c66e699708429
+size 14675

metadata/audioset/eval_lst.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8efd7d12ad5d6658b0027ec11fe084aa6e5790d0b8a41da0292ab0dce348daf7
+size 2019754

metadata/audioset/ontology.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c685f4403eecc3ca9be37fd7285cf212feaaea6ff7229d3e7ca89e0d1f2d15d
+size 342780

metadata/audioset/train_lst.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:230032344363fe8f23ac9e3d96ed54efd4a8f2477cd658efa4cbf78d13e5253f
+size 225874340

metadata/data_engine_json/audioset_iter2_train_data_engine_th10.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8cd5db4da07f7755e4c318aba538e87040c701e463883ce9e860652684c48da1
+size 329874536

metadata/data_engine_json/audioset_iter2_train_data_engine_th15.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bab017fe44d4964848a6a1ca67a1dc5fe0966b7ef4fa72c01336238a7c4dbb42
+size 242464513

metadata/data_engine_meta/child_label/bal_train_segments.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bf9959d7830a2e3122e6aec07e6d65c3282c9067b95700ebc79285324dbfb68
+size 1501593

metadata/data_engine_meta/child_label/bal_train_segments_multi_label.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e7cf8eb7237a8f156f963e0d39376d8e5f87b4b0fe386de64c1d527b9b0c20a
+size 1063491

metadata/data_engine_meta/child_label/eval_segments.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:503a99a3902e94f2569fecb801ca14da4526d525dbdee4cef08a4ec3ded77d26
+size 1449092

metadata/data_engine_meta/child_label/unbal_train_segments.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76ea2171235b11ce8803390b0b430a28f8c5ddaf70c9ee1bb3d3ceb0d6ced27a
+size 119776930

metadata/data_engine_meta/child_label/unbal_train_segments_multi_label.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf64f07ed2e96e67db83314bb4c29e98c117f12c95b3c7de431243f3d99aaec6
+size 66899715

metadata/evaluation/audiocaps_caption_eval.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b96a57ab8800d88cab2ecff108ca8e1cc5309b5505a83b04eb8fd956fcd4ef65
+size 741741

metadata/evaluation/audiocaps_label_eval.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8c2534f2f4c51d293d16a7116e5f8036601ec43adba2e7fd6682d8c5956c19b
+size 497487

metadata/evaluation/audiocaps_label_eval_sep_silence_test.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1463d450becd35d9dd1dd2769030da69de3e0c56cd4f78c67886228d04593c6b
+size 592858

metadata/evaluation/audioset_eval.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:74e77cc633789d8d798efffd8183d557b948f31426cb1a450f49aa6b43ec9378
+size 2411901

metadata/evaluation/audioset_single_eval.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e8faf73e0b84db5d49383f67dfacdd8aba491bb6faa580e4070a44c69c411646
+size 794337

metadata/evaluation/esc_eval.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1241a955f700332888bf40e80bc51cdaa6c9376edca09e5bdd05d5a7c25e32f0
+size 608128

metadata/evaluation/esc_eval_samples.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fe79e1aa7b60266379fd02556be36aa5e8aecaa9f6e6000e53e6c8d41f12c00
+size 11564

metadata/training/audiocaps_caption_32k_train.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fd4db95b685e9b747b842290e35c83301ed47db9be146a05d15a21de4d84709
+size 5660553

metadata/training/audiocaps_caption_32k_val.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2c1535eb967d344299013dc8d9e31cb657faf419f6924adf000ca6ebc3c32ac
+size 145573

metadata/training/audiocaps_label_32k_train_sep.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bc819bf9bf9e100c5f2553e5e8c49e4b9ef2e973c4b54bc4815a1f8fdbf7b9eb
+size 4869583

metadata/training/audiocaps_label_32k_val_sep.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:113de99477ccd9e2217e490ef2861f1f9197e6edcfc94ca451beabb36c824a16
+size 47563

metadata/training/audioset_32k_train_sep.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c713ab96a393ee913146ffe8cf2e4be77502868718ee2ac0d7feb23f0a636df6
+size 220302378

metadata/training/audioset_32k_train_sep_th10_iter2.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ede407a78b9d5e98157ca9a57e84aed887543affefa86c55642636b7f77ada52
+size 451618397

metadata/training/audioset_32k_train_sep_th15_iter2.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90df3df917d46b67c990fa506f26f97f8081b0e3879d460225f6438c2e65dada
+size 389681938

metadata/training/audioset_32k_val_sep.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66086b740b273f47b9248de8f94d9dd3fe94c6f790b9962f8c8e0ddd9ed5459f
+size 1473876

model/CLAPSep.py

@@ -0,0 +1,229 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main 
+@File    ：CLAPSep.py
+@IDE     ：PyCharm 
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2024/2/28 下午1:12 
+'''
+
+import torch
+import laion_clap
+from torchmetrics.audio.snr import(
+    scale_invariant_signal_noise_ratio as si_snr,
+    signal_noise_ratio as snr)
+from torchmetrics.audio.sdr import(
+    signal_distortion_ratio as sdr,
+    scale_invariant_signal_distortion_ratio as si_sdr)
+import copy
+import loralib as lora
+from torchlibrosa import ISTFT, STFT, SpecAugmentation
+from torchlibrosa.stft import magphase
+import librosa
+import pytorch_lightning as pl
+
+
+def loss_fn(pred, tgt):
+    return -0.9 * snr(pred, tgt).mean() - 0.1 * si_snr(pred, tgt).mean()
+
+
+def set_module(model, submodule_key, module):
+    tokens = submodule_key.split('.')
+    sub_tokens = tokens[:-1]
+    cur_mod = model
+    for s in sub_tokens:
+        cur_mod = getattr(cur_mod, s)
+    setattr(cur_mod, tokens[-1], module)
+
+
+def process_model(model, rank):
+    for n, module in model.named_modules():
+        if 'WindowAttention' in str(type(module)):
+            for n_, layer in module.named_modules():
+                if isinstance(layer, torch.nn.Linear):
+                    lora_layer = lora.Linear(layer.in_features, layer.out_features, r=rank,
+                                             bias=hasattr(layer, 'bias'), merge_weights=False)
+                    lora_layer.weight = layer.weight
+                    if hasattr(layer, 'bias'):
+                        lora_layer.bias = layer.bias
+                    set_module(model, n+'.'+n_, lora_layer)
+    return model
+
+
+class LightningModule(pl.LightningModule):
+    def __init__(self, clap_model, decoder_model, lr, use_lora=False, rank=8, nfft=1024):
+        super().__init__()
+        self.phase = decoder_model.phase
+        self.lr = lr
+        self.clap_model = clap_model
+        for p in self.clap_model.parameters():
+            p.requires_grad = False
+        self.audio_branch = copy.deepcopy(self.clap_model.model.audio_branch)
+        if use_lora:
+            process_model(self.audio_branch, rank)
+            lora.mark_only_lora_as_trainable(self.audio_branch, bias='lora_only')
+
+        self.decoder_model = decoder_model
+        self.stft = STFT(n_fft=nfft, hop_length=320,
+                         win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                         freeze_parameters=True)
+        self.istft = ISTFT(n_fft=nfft, hop_length=320,
+                           win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                           freeze_parameters=True)
+        self.features = self.install_forward_hooks()
+
+    def training_step(self, batch, batch_idx):
+        self.clap_model.eval()
+        self.audio_branch.eval()
+        # print([len(x) for x in batch])
+        mixed, mixed_resample, pos_cap, neg_cap, gt, pos_sample, neg_sample = batch
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        with torch.no_grad():
+            a = torch.rand((1,)).type_as(gt)
+            embed_pos_a, embed_neg_a = torch.chunk(
+                self.clap_model.get_audio_embedding_from_data(torch.concat([pos_sample, neg_sample], dim=0),
+                                                              use_tensor=True), dim=0, chunks=2)
+            embed_pos_t, embed_neg_t = torch.chunk(
+                self.clap_model.get_text_embedding(pos_cap + neg_cap, use_tensor=True), dim=0, chunks=2)
+            embed_pos = a * embed_pos_a + (1 - a) * embed_pos_t
+            embed_neg = a * embed_neg_a + (1 - a) * embed_neg_t
+        del self.features[:]
+        self.features.append(mag)
+        self.audio_branch({"waveform": mixed_resample})
+        a = torch.rand((1,))
+        if a < 0.25:
+            loss = self.cal_loss(embed_pos, torch.zeros_like(embed_pos), mag, cos, sin, length=mixed.size(-1), gt=gt)
+        elif a < 0.5:
+            loss = self.cal_loss(torch.zeros_like(embed_neg), embed_neg, mag, cos, sin, length=mixed.size(-1), gt=gt)
+        else:
+            loss = self.cal_loss(embed_pos, embed_neg, mag, cos, sin, length=mixed.size(-1), gt=gt)
+        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, sync_dist=True, batch_size=len(mixed))
+        del self.features[:]
+        return loss
+
+    def cal_loss(self, embed_p, embed_n, mag, cos, sin, length, gt):
+        embed = torch.nn.functional.normalize(torch.concat([embed_p, embed_n], dim=-1), dim=-1)
+        mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+        pred = self.wav_reconstruct(mask, mag, cos, sin, length=length)
+        return loss_fn(pred, gt)
+
+    def wav_reconstruct(self, mask, mag_x, cos_x, sin_x, length):
+        # ref: https://github.com/Audio-AGI/AudioSep/blob/main/models/resunet.py
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        if self.phase:
+            mag_y = torch.nn.functional.relu_(mag_x * mask[0])
+            _, mask_cos, mask_sin = magphase(mask[1], mask[2])
+            cos_y = cos_x * mask_cos - sin_x * mask_sin
+            sin_y = sin_x * mask_cos + cos_x * mask_sin
+        else:
+            mag_y = torch.nn.functional.relu_(mag_x * mask)
+            cos_y = cos_x
+            sin_y = sin_x
+        pred = self.istft(mag_y * cos_y, mag_y * sin_y, length=length)
+        return pred
+
+    def validation_step(self, batch, batch_idx):
+        mixed, mixed_resample, label, neg_label, gt, _, _ = batch
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        self.features.append(mag)
+        with torch.no_grad():
+            embed_pos = self.clap_model.get_text_embedding(label, use_tensor=True)
+            embed_neg = self.clap_model.get_text_embedding(neg_label, use_tensor=True)
+            embed = torch.concat([embed_pos, embed_neg], dim=-1)
+            self.audio_branch({"waveform": mixed_resample})
+            mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+            pred = self.wav_reconstruct(mask, mag, cos, sin, length=mixed.size(-1))
+            loss = si_snr(pred, gt).mean() - si_snr(mixed, gt).mean()
+        del self.features[:]
+        self.log("val_loss", loss, on_epoch=True, prog_bar=True, sync_dist=True, batch_size=len(mixed))
+        return {"val_loss": loss}
+
+    def on_test_start(self) -> None:
+        self.sdr_vals = torch.tensor([])
+        self.sdri_vals = torch.tensor([])
+        self.sisdr_vals = torch.tensor([])
+        self.sisdri_vals = torch.tensor([])
+
+    def test_step(self, batch, batch_idx):
+        mixed, mixed_resample, label, neg_label, gt = batch
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        with torch.no_grad():
+            embed_pos_bached, embed_neg_bached = torch.chunk(self.clap_model.get_text_embedding(label + neg_label, use_tensor=True), chunks=2, dim=0)
+            del self.features[:]
+            # only positive
+            # embed = torch.concat([embed_pos_bached, torch.zeros_like(embed_neg_bached)], dim=1)
+            # only negative
+            # embed = torch.concat([torch.zeros_like(embed_pos_bached), embed_neg_bached], dim=1)
+            # positive and negative
+            embed = torch.concat([embed_pos_bached, embed_neg_bached], dim=1)
+            self.features.append(mag)
+            self.audio_branch({"waveform": mixed_resample})
+            mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+            pred = self.wav_reconstruct(mask, mag, cos, sin, length=mixed.size(-1))
+            sisdr = si_sdr(pred, gt).cpu()
+            self.sisdr_vals = torch.concat([self.sisdr_vals, sisdr])
+            self.sisdri_vals = torch.concat([self.sisdri_vals, sisdr - si_sdr(mixed, gt).cpu()])
+            sdr_ = sdr(pred, gt).cpu()
+            self.sdr_vals = torch.concat([self.sdr_vals, sdr_])
+            self.sdri_vals = torch.concat([self.sdri_vals, sdr_ - sdr(mixed, gt).cpu()])
+        del self.features[:]
+    
+    def on_test_end(self) -> None:
+        print(f"SDR-mean: {torch.mean(self.sdr_vals).cpu().numpy():.4f}, SDR-std: {torch.std(self.sdr_vals).cpu().numpy():.4f}")
+        print(f"SDRi-mean: {torch.mean(self.sdri_vals).cpu().numpy():.4f}, SDRi-std: {torch.std(self.sdri_vals).cpu().numpy():.4f}")
+        print(f"SISDR-mean: {torch.mean(self.sisdr_vals).cpu().numpy():.4f}, SISDR-std: {torch.std(self.sisdr_vals).cpu().numpy():.4f}")
+        print(f"SISDRi-mean: {torch.mean(self.sisdri_vals).cpu().numpy():.4f}, SISDRi-std: {torch.std(self.sisdri_vals).cpu().numpy():.4f}")
+    
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
+        schedular = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.3, patience=5,
+                                                                        verbose=True, min_lr=5e-6)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": schedular,
+                "interval": "epoch",
+                "monitor": "val_loss"
+            },
+        }
+
+    def install_forward_hooks(self):
+        features = []
+        spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
+                                          freq_drop_width=8, freq_stripes_num=2)
+
+        def get_features_list(_, __, output):
+            features.append(output)
+
+        def get_features_list_basic_layer(_, __, output):
+            features.append(output[0])
+
+        def spec_augmentation_hook(_, __, out):
+            out = out.transpose(1, 3)
+            out = spec_augmenter(out)
+            return out.transpose(1, 3)
+
+        def spectrogram_padding(_, __, out):
+            return torch.nn.functional.pad(out, (0, 0, 0, 1024 - out.size(2)))
+
+        self.clap_model.model.audio_branch.bn0.register_forward_hook(spec_augmentation_hook)
+        self.audio_branch.spectrogram_extractor.register_forward_hook(spectrogram_padding)
+        self.audio_branch.patch_embed.register_forward_hook(get_features_list)
+        for module in self.audio_branch.layers:
+            module.register_forward_hook(get_features_list_basic_layer)
+        return features
+
+    # # this will only save tuned parameters during training
+    # def on_save_checkpoint(self, checkpoint):
+    #     weights = checkpoint['state_dict']
+    #     new_dict = {}
+    #     for k, v in weights.items():
+    #         if any(e in k for e in ['lora', 'attn.qkv.bias', 'attn.proj.bias', 'decoder_model']):
+    #             new_dict[k] = v
+    #     checkpoint['state_dict'] = new_dict

model/CLAPSep_decoder.py

@@ -0,0 +1,605 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main
+@File    ：CLAPsep_decoder.py
+@IDE     ：PyCharm
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2023/10/31 下午8:34
+'''
+
+from laion_clap.clap_module.htsat import *
+from einops import rearrange
+import numpy as np
+
+class Transpose(nn.Module):
+
+    def __init__(self, dim0, dim1):
+        super(Transpose, self).__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x):
+        return x.transpose(self.dim0, self.dim1)
+
+
+class Swish(nn.Module):
+
+    def __init__(self):
+        super(Swish, self).__init__()
+
+    def forward(self, x):
+        return x * x.sigmoid()
+
+
+class Glu(nn.Module):
+
+    def __init__(self, dim):
+        super(Glu, self).__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        x_in, x_gate = x.chunk(2, dim=self.dim)
+        return x_in * x_gate.sigmoid()
+
+
+class FiLM(nn.Module):
+    def __init__(self, dim_in=1024, hidden_dim=768):
+        super(FiLM, self).__init__()
+        self.beta = nn.Linear(dim_in, hidden_dim)
+        self.gamma = nn.Linear(dim_in, hidden_dim)
+
+    def forward(self, hidden_state, embed):
+        embed = embed.unsqueeze(1)
+        return self.gamma(embed) * hidden_state + self.beta(embed)
+
+
+class SkipTrans(nn.Module):
+    def __init__(self, in_features, out_features, embed_dim=512, film=True):
+        super(SkipTrans, self).__init__()
+        self.film = film
+        if film:
+            self.skip_conv = FiLM(embed_dim, out_features)
+        self.feature_proj = nn.Linear(in_features, out_features)
+        self.norm = nn.LayerNorm(out_features)
+
+    def forward(self, skip, embed, x=None):
+        out = self.feature_proj(skip)
+        if self.film:
+            out = self.skip_conv(out, embed)
+        return self.norm(out) if x is None else self.norm(out + x)
+
+class Conv1d(nn.Conv1d):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride = 1,
+        padding = "same",
+        dilation = 1,
+        groups = 1,
+        bias = True
+    ):
+        super(Conv1d, self).__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=0,
+            dilation=dilation,
+            groups=groups,
+            bias=bias,
+            padding_mode="zeros")
+
+        # Assert
+        assert padding in ["valid", "same", "causal"]
+
+        # Padding
+        if padding == "valid":
+            self.pre_padding = None
+        elif padding == "same":
+            self.pre_padding = nn.ConstantPad1d(padding=((kernel_size - 1) // 2, (kernel_size - 1) // 2), value=0)
+        elif padding == "causal":
+            self.pre_padding = nn.ConstantPad1d(padding=(kernel_size - 1, 0), value=0)
+
+        # Variational Noise
+        self.noise = None
+        self.vn_std = None
+
+    def init_vn(self, vn_std):
+
+        # Variational Noise
+        self.vn_std = vn_std
+
+    def sample_synaptic_noise(self, distributed):
+
+        # Sample Noise
+        self.noise = torch.normal(mean=0.0, std=1.0, size=self.weight.size(), device=self.weight.device, dtype=self.weight.dtype)
+
+        # Broadcast Noise
+        if distributed:
+            torch.distributed.broadcast(self.noise, 0)
+
+    def forward(self, input):
+
+        # Weight
+        weight = self.weight
+
+        # Add Noise
+        if self.noise is not None and self.training:
+            weight = weight + self.vn_std * self.noise
+
+        # Padding
+        if self.pre_padding is not None:
+            input = self.pre_padding(input)
+
+        # Apply Weight
+        return F.conv1d(input, weight, self.bias, self.stride, self.padding, self.dilation, self.groups)
+
+
+class ConvolutionModule(nn.Module):
+    """Conformer Convolution Module
+
+    Args:
+        dim_model: input feature dimension
+        dim_expand: output feature dimension
+        kernel_size: 1D depthwise convolution kernel size
+        Pdrop: residual dropout probability
+        stride: 1D depthwise convolution stride
+        padding: "valid", "same" or "causal"
+
+    Input: (batch size, input length, dim_model)
+    Output: (batch size, output length, dim_expand)
+
+    """
+
+    def __init__(self, dim_model, dim_expand, kernel_size, Pdrop, stride, padding):
+        super(ConvolutionModule, self).__init__()
+
+        # Layers
+        self.layers = nn.Sequential(
+            nn.LayerNorm(dim_model, eps=1e-6),
+            Transpose(1, 2),
+            Conv1d(dim_model, 2 * dim_expand, kernel_size=1),
+            Glu(dim=1),
+            Conv1d(dim_expand, dim_expand, kernel_size, stride=stride, padding=padding, groups=dim_expand),
+            nn.BatchNorm1d(dim_expand),
+            Swish(),
+            Conv1d(dim_expand, dim_expand, kernel_size=1),
+            Transpose(1, 2),
+            nn.Dropout(p=Pdrop)
+        )
+        self.ln = nn.LayerNorm(dim_expand)
+
+    def forward(self, x):
+        return self.ln(self.layers(x)+x)
+
+
+class BasicLayerDec(nn.Module):
+    """ A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Local window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
+    """
+
+    def __init__(self, dim, input_resolution, depth, num_heads, window_size,
+                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,
+                 norm_before_mlp='ln'):
+
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
+                                 num_heads=num_heads, window_size=window_size,
+                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
+                                 mlp_ratio=mlp_ratio,
+                                 qkv_bias=qkv_bias, qk_scale=qk_scale,
+                                 drop=drop, attn_drop=attn_drop,
+                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
+                                 norm_layer=norm_layer, norm_before_mlp=norm_before_mlp)
+            for i in range(depth)])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample((input_resolution[0]//2, input_resolution[1]//2), dim=dim * 2, norm_layer=norm_layer)
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        attns = []
+        if self.downsample is not None:
+            x = self.downsample(x)
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x, attn = blk(x)
+                if not self.training:
+                    attns.append(attn.unsqueeze(0))
+        if not self.training:
+            attn = torch.cat(attns, dim = 0)
+            attn = torch.mean(attn, dim = 0)
+        return x, attn
+
+    def extra_repr(self):
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+
+class PatchExpand(nn.Module):
+    def __init__(self, input_resolution, dim, dim_scale=2, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.expand = nn.Linear(dim, 2 * dim, bias=False) if dim_scale == 2 else nn.Identity()
+        self.norm = norm_layer(dim // dim_scale)
+
+    def forward(self, x):
+        """
+        x: B, H*W, C
+        """
+        H, W = self.input_resolution
+        x = self.expand(x)
+        B, L, C = x.shape
+        assert L == H * W, "input feature has wrong size"
+
+        x = x.view(B, H, W, C)
+        # This is the original implementation in SwinUnet
+        # x = rearrange(x, 'b h w (p1 p2 c)-> b (h p1) (w p2) c', p1=2, p2=2, c=C // 4)
+
+        # here is our implementation
+        # can reverse patch-merging in Swin-Transformer encoder, seems helpful
+        x0, x2, x1, x3 = x.chunk(4, dim=-1)
+        x = torch.stack((x0, x1, x2, x3), dim=-1)
+        x = torch.chunk(x, C // 4, dim=-2)
+        x = torch.concat(x, dim=-1).squeeze(-2)
+        x = rearrange(x, 'b h w c -> b c h w')
+        x = torch.nn.functional.pixel_shuffle(x, 2)
+        x = rearrange(x, 'b c h w -> b h w c')
+        x = x.view(B, -1, C // 4)
+        x = self.norm(x)
+
+        return x
+
+
+class InversePatchEmbed(nn.Module):
+    """
+    Patch Embedding to 2D Image.
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True,
+                 patch_stride=16):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patch_stride = to_2tuple(patch_stride)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patch_stride = patch_stride
+        self.grid_size = (img_size[0] // patch_stride[0], img_size[1] // patch_stride[1])
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)
+
+        self.proj = nn.ConvTranspose2d(embed_dim, in_chans, kernel_size=patch_size, stride=patch_stride, padding=padding)
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x):
+        # B, C, H, W = x.shape
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.norm(x)
+        if self.flatten:
+            # x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+            x = x.transpose(1, 2).unflatten(2, self.grid_size).contiguous()  # BNC -> BCHW
+        x = self.proj(x)
+
+        return x
+
+
+class HTSAT_Decoder(nn.Module):
+    r"""HTSAT_decoder based on the Swin Transformer
+    Args:
+        spec_size (int | tuple(int)): Input Spectrogram size. Default 256
+        patch_size (int | tuple(int)): Patch size. Default: 4
+        path_stride (iot | tuple(int)): Patch Stride for Frequency and Time Axis. Default: 4
+        in_chans (int): Number of input image channels. Default: 1 (mono)
+        num_classes (int): Number of classes for classification head. Default: 527
+        embed_dim (int): Patch embedding dimension. Default: 96
+        depths (tuple(int)): Depth of each HTSAT-Swin Transformer layer.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size. Default: 8
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
+        drop_rate (float): Dropout rate. Default: 0
+        attn_drop_rate (float): Attention dropout rate. Default: 0
+        drop_path_rate (float): Stochastic depth rate. Default: 0.1
+        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
+        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
+        patch_norm (bool): If True, add normalization after patch embedding. Default: True
+        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
+    """
+
+    def __init__(self, lan_embed_dim=512, spec_size=256, patch_size=4, patch_stride=(4, 4),
+                 in_chans=1, num_classes=527,
+                 embed_dim=48, depths=[1, 1, 1, 1], num_heads=[4, 8, 16, 32],
+                 window_size=8, mlp_ratio=4., qkv_bias=True, qk_scale=None,
+                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
+                 norm_layer=nn.LayerNorm,
+                 ape=False, patch_norm=True,
+                 use_checkpoint=False, norm_before_mlp='ln', encoder_embed_dim=96, phase=False,
+                 spec_factor=8, d_attn=640, n_masker_layer=4, conv=False):
+        super(HTSAT_Decoder, self).__init__()
+        self.mel_bins = 64
+        self.spec_size = spec_size
+        self.phase = phase
+        self.patch_stride = patch_stride
+        self.patch_size = patch_size
+        self.window_size = window_size
+        self.embed_dim = embed_dim
+        self.depths = depths
+        self.ape = ape
+        self.in_chans = in_chans
+        self.num_classes = num_classes
+        self.num_heads = num_heads
+        self.num_layers = len(self.depths)
+        self.num_features = int(self.embed_dim * 2 ** (self.num_layers - 1))
+
+        self.drop_rate = drop_rate
+        self.attn_drop_rate = attn_drop_rate
+        self.drop_path_rate = drop_path_rate
+
+        self.qkv_bias = qkv_bias
+        self.qk_scale = None
+
+        self.patch_norm = patch_norm
+        self.norm_layer = norm_layer if self.patch_norm else None
+        self.norm_before_mlp = norm_before_mlp
+        self.mlp_ratio = mlp_ratio
+
+        self.use_checkpoint = use_checkpoint
+
+        #  process mel-spec ; used only once
+        self.freq_ratio = self.spec_size // self.mel_bins
+
+
+        # split spctrogram into non-overlapping patches
+        self.inverse_patch_embed = InversePatchEmbed(
+            img_size=self.spec_size, patch_size=self.patch_size, in_chans=self.in_chans,
+            embed_dim=self.embed_dim, norm_layer=self.norm_layer, patch_stride=patch_stride)
+
+        patches_resolution = self.inverse_patch_embed.grid_size
+        self.patches_resolution = patches_resolution
+
+
+        # stochastic depth
+        dpr = [x.item() for x in
+               torch.linspace(0, self.drop_path_rate, sum(self.depths))]  # stochastic depth decay rule
+
+        # build layers
+        self.layers = nn.ModuleList()
+        self.skip = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayerDec(dim=int(self.embed_dim * 2 ** i_layer),
+                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
+                                                 patches_resolution[1] // (2 ** i_layer)),
+                               depth=self.depths[i_layer],
+                               num_heads=self.num_heads[i_layer],
+                               window_size=self.window_size,
+                               mlp_ratio=self.mlp_ratio,
+                               qkv_bias=self.qkv_bias, qk_scale=self.qk_scale,
+                               drop=self.drop_rate, attn_drop=self.attn_drop_rate,
+                               drop_path=dpr[sum(self.depths[:i_layer]):sum(self.depths[:i_layer + 1])],
+                               norm_layer=self.norm_layer,
+                               downsample=PatchExpand if (i_layer < self.num_layers - 1) else None,
+                               use_checkpoint=use_checkpoint,
+                               norm_before_mlp=self.norm_before_mlp)
+            self.layers.append(layer)
+            self.skip.append(
+                SkipTrans(embed_dim=lan_embed_dim, in_features=int(encoder_embed_dim * 2 ** i_layer), out_features=int(self.embed_dim * 2 ** i_layer)),
+            )
+        self.layers = self.layers[::-1]
+        self.skip = self.skip[::-1]
+        # self.skip.append(
+        #     SkipTrans(embed_dim=lan_embed_dim, in_features=self.mel_bins, out_features=self.mel_bins),
+        # )
+
+        d_spec = self.mel_bins * spec_factor + 1
+
+        self.spec_norm = nn.BatchNorm2d(d_spec, momentum=0.01)
+        self.conv = conv
+        if not conv:
+            encoder_layer = nn.TransformerEncoderLayer(d_model=d_attn, nhead=8,
+                                                       dim_feedforward=int(d_attn * self.mlp_ratio),
+                                                       batch_first=True, dropout=0)
+            transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_masker_layer)
+
+            self.mask_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_attn),
+                nn.LayerNorm(d_attn),
+                transformer_encoder,
+                nn.Linear(d_attn, d_spec)
+            )
+        else:
+            self.mask_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_spec),
+                nn.LayerNorm(d_spec),
+                *[ConvolutionModule(dim_model=d_spec, dim_expand=d_spec, kernel_size=9, padding='same',
+                                    Pdrop=0, stride=1) for i in range(n_masker_layer)]
+            )
+        if self.phase:
+            self.phase_net = nn.Sequential(
+                nn.Linear(self.mel_bins + d_spec, d_spec * 2),
+                nn.LayerNorm(d_spec * 2),
+                *[ConvolutionModule(dim_model=d_spec * 2, dim_expand=d_spec * 2, kernel_size=9, padding='same',
+                                    Pdrop=0, stride=1) for i in range(n_masker_layer)]
+            )
+
+        self.film = SkipTrans(embed_dim=lan_embed_dim, in_features=encoder_embed_dim * 8, out_features=self.num_features)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    # @torch.jit.ignore
+    # def no_weight_decay(self):
+    #     return {'absolute_pos_embed'}
+    #
+    # @torch.jit.ignore
+    # def no_weight_decay_keywords(self):
+    #     return {'relative_position_bias_table'}
+
+    def forward(self, hidden_state, skip_features, embed):
+        skip_features = skip_features[::-1]
+        # hidden_state = torch.randn(hidden_state.shape).type_as(hidden_state)
+
+        spec = skip_features[-1]
+
+        h = self.film(hidden_state, embed)
+
+        for i, (layer, f, skip) in enumerate(zip(self.layers, skip_features, self.skip)):
+            h = layer(h)[0]
+            h = skip(skip=f, embed=embed, x=h)
+
+        h = self.reshape_img2wav(self.inverse_patch_embed(h)).squeeze(1)
+
+        h = h[:, :spec.size(2), :]
+
+        spec = spec.transpose(1, 3)
+
+        spec = self.spec_norm(spec).transpose(1, 3).squeeze(1)
+
+        h = torch.concat([spec, h], dim=-1)
+
+        mask = self.mask_net(h).unsqueeze(1)
+
+        if self.phase:
+            mask_r, mask_i = torch.chunk(self.phase_net(h).unsqueeze(1), chunks=2, dim=-1)
+            return torch.sigmoid(mask), torch.tanh(mask_r), torch.tanh(mask_i)
+        else:
+            return torch.sigmoid(mask)
+
+    def reshape_img2wav(self, x):
+        # (B, 1, 256, 256)
+        x = x.reshape(x.shape[0], x.shape[1], self.freq_ratio, x.shape[2]//self.freq_ratio, x.shape[3]) # (B, 1, 4, 64, 256)
+        x = x.permute(0, 1, 3, 2, 4).contiguous()
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2], x.shape[3] * x.shape[4])
+        x = x.permute(0, 1, 3, 2).contiguous()
+        return x
+
+
+# if __name__ == "__main__":
+#     import torch
+#     from msclap import CLAP
+#     import os
+#     import torchaudio
+#     import torchaudio.transforms as T
+#     import numpy as np
+#     import random
+#     from torchlibrosa import Spectrogram, LogmelFilterBank
+#     clap_model = CLAP(model_fp="/home/user/202212661/clapsep/Waveformer-main/checkpoint_path/CLAP_weights_2023.pth",
+#                       version='2023', use_cuda=True)
+#     text_data = [
+#         "Acoustic_guitar", "Applause", "Bark", "Bass_drum", "Burping_or_eructation",
+#         "Bus", "Cello", "Chime", "Clarinet", "Computer_keyboard",
+#         "Cough", "Cowbell", "Double_bass", "Drawer_open_or_close", "Electric_piano",
+#         "Fart", "Finger_snapping", "Fireworks", "Flute", "Glockenspiel",
+#         "Gong", "Gunshot_or_gunfire", "Harmonica", "Hi-hat", "Keys_jangling",
+#         "Knock", "Laughter", "Meow", "Microwave_oven", "Oboe",
+#         "Saxophone", "Scissors", "Shatter", "Snare_drum", "Squeak",
+#         "Tambourine", "Tearing", "Telephone", "Trumpet", "Violin_or_fiddle",
+#         "Writing"]
+#     # Extract text embeddings
+#     text_embeddings = clap_model.get_text_embeddings(text_data)
+#     path = "/home/user/202212661/clapsep/Waveformer-main/data/FSDSoundScapes/FSDKaggle2018/train/Tearing/2232ce13.wav"
+#     # Extract audio embeddings
+#     audio_embeddings_ = clap_model.get_audio_embeddings([path])
+#
+#     window = 'hann'
+#     center = True
+#     pad_mode = 'reflect'
+#     ref = 1.0
+#     amin = 1e-10
+#     top_db = None
+#
+#     spectrogram_extractor = Spectrogram(n_fft=512, hop_length=160,
+#                                         win_length=512, window=window, center=center, pad_mode=pad_mode,
+#                                         freeze_parameters=True).cuda()
+#     # Logmel feature extractor
+#     logmel_extractor = LogmelFilterBank(sr=16000, n_fft=512,
+#                                         n_mels=64, fmin=0, fmax=8000, ref=ref, amin=amin,
+#                                         top_db=top_db,
+#                                         freeze_parameters=True).cuda()
+#
+#     clap_model.clap.audio_encoder.base.htsat.spectrogram_extractor = spectrogram_extractor
+#     clap_model.clap.audio_encoder.base.htsat.logmel_extractor = logmel_extractor
+#
+#     features = []
+#
+#
+#     def get_features_list(module, input, output):
+#         features.append(output)
+#
+#
+#     def get_features_list_basic_layer(module, input, output):
+#         features.append(output[0])
+#
+#
+#     clap_model.clap.audio_encoder.base.htsat.patch_embed.register_forward_hook(get_features_list)
+#     for module in clap_model.clap.audio_encoder.base.htsat.layers:
+#         module.register_forward_hook(get_features_list_basic_layer)
+#
+#     audio_time_series, sample_rate = torchaudio.load(path)
+#     resample_rate = 16000
+#     if resample_rate != sample_rate:
+#         resampler = T.Resample(sample_rate, resample_rate)
+#         audio_time_series = resampler(audio_time_series)
+#
+#     sample_rate = resample_rate
+#     audio_duration = 10
+#     audio_time_series = audio_time_series.reshape(-1)
+#     if audio_duration * sample_rate >= audio_time_series.shape[0]:
+#         repeat_factor = int(np.ceil((audio_duration * sample_rate) /
+#                                     audio_time_series.shape[0]))
+#         # Repeat audio_time_series by repeat_factor to match audio_duration
+#         audio_time_series = audio_time_series.repeat(repeat_factor)
+#         # remove excess part of audio_time_series
+#         audio_time_series = audio_time_series[0:audio_duration * sample_rate]
+#     else:
+#         # audio_time_series is longer than predefined audio duration,
+#         # so audio_time_series is trimmed
+#         start_index = random.randrange(
+#             audio_time_series.shape[0] - audio_duration * sample_rate)
+#         audio_time_series = audio_time_series[start_index:start_index +
+#                                                           audio_duration * sample_rate]
+#

model/CLAPSep_infer.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python
+# -*- coding: UTF-8 -*-
+'''
+@Project ：Waveformer-main 
+@File    ：CLAPSep.py
+@IDE     ：PyCharm 
+@Author  ：Aisaka/Hao Ma @SDU
+@Date    ：2024/2/28 下午1:12 
+'''
+
+import os
+import torch
+import laion_clap
+from torchmetrics.audio.snr import(
+    scale_invariant_signal_noise_ratio as si_snr,
+    signal_noise_ratio as snr)
+from torchmetrics.audio.sdr import(
+    signal_distortion_ratio as sdr,
+    scale_invariant_signal_distortion_ratio as si_sdr)
+import copy
+import loralib as lora
+from torchlibrosa import ISTFT, STFT, SpecAugmentation
+from torchlibrosa.stft import magphase
+import librosa
+import pytorch_lightning as pl
+import soundfile as sf
+
+
+def loss_fn(pred, tgt):
+    return -0.9 * snr(pred, tgt).mean() - 0.1 * si_snr(pred, tgt).mean()
+
+
+def set_module(model, submodule_key, module):
+    tokens = submodule_key.split('.')
+    sub_tokens = tokens[:-1]
+    cur_mod = model
+    for s in sub_tokens:
+        cur_mod = getattr(cur_mod, s)
+    setattr(cur_mod, tokens[-1], module)
+
+
+def process_model(model, rank):
+    for n, module in model.named_modules():
+        if 'WindowAttention' in str(type(module)):
+            for n_, layer in module.named_modules():
+                if isinstance(layer, torch.nn.Linear):
+                    lora_layer = lora.Linear(layer.in_features, layer.out_features, r=rank,
+                                             bias=hasattr(layer, 'bias'), merge_weights=False)
+                    lora_layer.weight = layer.weight
+                    if hasattr(layer, 'bias'):
+                        lora_layer.bias = layer.bias
+                    set_module(model, n+'.'+n_, lora_layer)
+    return model
+
+
+class LightningModule(pl.LightningModule):
+    def __init__(self, clap_model, decoder_model, lr, use_lora=False, rank=8, nfft=1024):
+        super().__init__()
+        self.phase = decoder_model.phase
+        self.lr = lr
+        self.clap_model = clap_model
+        for p in self.clap_model.parameters():
+            p.requires_grad = False
+        self.audio_branch = copy.deepcopy(self.clap_model.model.audio_branch)
+        if use_lora:
+            process_model(self.audio_branch, rank)
+            lora.mark_only_lora_as_trainable(self.audio_branch, bias='lora_only')
+
+        self.decoder_model = decoder_model
+        self.stft = STFT(n_fft=nfft, hop_length=320,
+                         win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                         freeze_parameters=True)
+        self.istft = ISTFT(n_fft=nfft, hop_length=320,
+                           win_length=nfft, window='hann', center=True, pad_mode='reflect',
+                           freeze_parameters=True)
+        self.features = self.install_forward_hooks()
+
+    def training_step(self, batch, batch_idx):
+        self.clap_model.eval()
+        self.audio_branch.eval()
+        # print([len(x) for x in batch])
+        mixed, mixed_resample, pos_cap, neg_cap, gt, pos_sample, neg_sample = batch
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        with torch.no_grad():
+            a = torch.rand((1,)).type_as(gt)
+            embed_pos_a, embed_neg_a = torch.chunk(
+                self.clap_model.get_audio_embedding_from_data(torch.concat([pos_sample, neg_sample], dim=0),
+                                                              use_tensor=True), dim=0, chunks=2)
+            embed_pos_t, embed_neg_t = torch.chunk(
+                self.clap_model.get_text_embedding(pos_cap + neg_cap, use_tensor=True), dim=0, chunks=2)
+            embed_pos = a * embed_pos_a + (1 - a) * embed_pos_t
+            embed_neg = a * embed_neg_a + (1 - a) * embed_neg_t
+        del self.features[:]
+        self.features.append(mag)
+        self.audio_branch({"waveform": mixed_resample})
+        a = torch.rand((1,))
+        if a < 0.25:
+            loss = self.cal_loss(embed_pos, torch.zeros_like(embed_pos), mag, cos, sin, length=mixed.size(-1), gt=gt)
+        elif a < 0.5:
+            loss = self.cal_loss(torch.zeros_like(embed_neg), embed_neg, mag, cos, sin, length=mixed.size(-1), gt=gt)
+        else:
+            loss = self.cal_loss(embed_pos, embed_neg, mag, cos, sin, length=mixed.size(-1), gt=gt)
+        self.log("train_loss", loss.item(), on_epoch=True, prog_bar=True, sync_dist=True, batch_size=len(mixed))
+        del self.features[:]
+        return loss
+
+    def cal_loss(self, embed_p, embed_n, mag, cos, sin, length, gt):
+        embed = torch.nn.functional.normalize(torch.concat([embed_p, embed_n], dim=-1), dim=-1)
+        mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+        pred = self.wav_reconstruct(mask, mag, cos, sin, length=length)
+        return loss_fn(pred, gt)
+
+    def wav_reconstruct(self, mask, mag_x, cos_x, sin_x, length):
+        # ref: https://github.com/Audio-AGI/AudioSep/blob/main/models/resunet.py
+        # Y = |Y|cos∠Y + j|Y|sin∠Y
+        #   = |Y|cos(∠X + ∠M) + j|Y|sin(∠X + ∠M)
+        #   = |Y|(cos∠X cos∠M - sin∠X sin∠M) + j|Y|(sin∠X cos∠M + cos∠X sin∠M)
+        if self.phase:
+            mag_y = torch.nn.functional.relu_(mag_x * mask[0])
+            _, mask_cos, mask_sin = magphase(mask[1], mask[2])
+            cos_y = cos_x * mask_cos - sin_x * mask_sin
+            sin_y = sin_x * mask_cos + cos_x * mask_sin
+        else:
+            mag_y = torch.nn.functional.relu_(mag_x * mask)
+            cos_y = cos_x
+            sin_y = sin_x
+        pred = self.istft(mag_y * cos_y, mag_y * sin_y, length=length)
+        return pred
+
+    def validation_step(self, batch, batch_idx):
+        mixed, mixed_resample, label, neg_label, gt, _, _ = batch
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        self.features.append(mag)
+        with torch.no_grad():
+            embed_pos = self.clap_model.get_text_embedding(label, use_tensor=True)
+            embed_neg = self.clap_model.get_text_embedding(neg_label, use_tensor=True)
+            embed = torch.concat([embed_pos, embed_neg], dim=-1)
+            self.audio_branch({"waveform": mixed_resample})
+            mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+            pred = self.wav_reconstruct(mask, mag, cos, sin, length=mixed.size(-1))
+            loss = si_snr(pred, gt).mean() - si_snr(mixed, gt).mean()
+        del self.features[:]
+        self.log("val_loss", loss, on_epoch=True, prog_bar=True, sync_dist=True, batch_size=len(mixed))
+        return {"val_loss": loss}
+
+
+    def test_step(self, batch, batch_idx):
+        mixed, mixed_resample, labels, src_path = batch
+        assert len(labels) == 1
+        src_path = src_path[0]
+        save_path = src_path.replace('/data/sound/audioset/audios_32k', 
+                                     'data_engine_infer/audioset_separation_child_label')[:-4]
+        os.makedirs(save_path, exist_ok=True)
+        real, imag = self.stft(mixed)
+        mag, cos, sin = magphase(real, imag)
+        with torch.no_grad():
+            labels = labels[0].split('|')
+            for pos_label in labels:
+                embed_pos = self.clap_model.get_text_embedding(pos_label, use_tensor=True)
+                
+                neg_labels = copy.deepcopy(labels)
+                neg_labels.remove(pos_label)
+                if neg_labels:
+                    neg_label = ', '.join(neg_labels)
+                    embed_neg = self.clap_model.get_text_embedding(neg_label, use_tensor=True)
+                else:
+                    embed_neg = torch.zeros_like(embed_pos)
+                # only positive
+                # embed = torch.concat([embed_pos, torch.zeros_like(embed_pos)], dim=1)
+                # positive and negative
+                embed = torch.concat([embed_pos, embed_neg], dim=1)
+                del self.features[:]
+                self.features.append(mag)
+                self.audio_branch({"waveform": mixed_resample})
+                mask = self.decoder_model(hidden_state=self.features[-1], skip_features=self.features[:-1], embed=embed)
+                pred = self.wav_reconstruct(mask, mag, cos, sin, length=mixed.size(-1))
+                del self.features[:]
+                sf.write(os.path.join(save_path, pos_label + '.wav'), pred.squeeze().cpu().numpy(), samplerate=32000)
+
+    
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.lr)
+        schedular = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.3, patience=5,
+                                                                        verbose=True, min_lr=5e-6)
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": schedular,
+                "interval": "epoch",
+                "monitor": "val_loss"
+            },
+        }
+
+    def install_forward_hooks(self):
+        features = []
+        spec_augmenter = SpecAugmentation(time_drop_width=64, time_stripes_num=2,
+                                          freq_drop_width=8, freq_stripes_num=2)
+
+        def get_features_list(_, __, output):
+            features.append(output)
+
+        def get_features_list_basic_layer(_, __, output):
+            features.append(output[0])
+
+        def spec_augmentation_hook(_, __, out):
+            out = out.transpose(1, 3)
+            out = spec_augmenter(out)
+            return out.transpose(1, 3)
+
+        def spectrogram_padding(_, __, out):
+            return torch.nn.functional.pad(out, (0, 0, 0, 1024 - out.size(2)))
+
+        self.clap_model.model.audio_branch.bn0.register_forward_hook(spec_augmentation_hook)
+        self.audio_branch.spectrogram_extractor.register_forward_hook(spectrogram_padding)
+        self.audio_branch.patch_embed.register_forward_hook(get_features_list)
+        for module in self.audio_branch.layers:
+            module.register_forward_hook(get_features_list_basic_layer)
+        return features
+
+    # # this will only save tuned parameters during training
+    # def on_save_checkpoint(self, checkpoint):
+    #     weights = checkpoint['state_dict']
+    #     new_dict = {}
+    #     for k, v in weights.items():
+    #         if any(e in k for e in ['lora', 'attn.qkv.bias', 'attn.proj.bias', 'decoder_model']):
+    #             new_dict[k] = v
+    #     checkpoint['state_dict'] = new_dict
+
+
+

requirements.txt

@@ -0,0 +1,13 @@
+# torch
+librosa
+# torchaudio
+# torchvision
+pytorch-lightning==2.4.0
+tensorboard==2.18.0
+torchlibrosa
+numpy==1.26.4
+einops
+loralib
+transformers==4.30.2
+laion-clap
+matplotlib

run.sh

@@ -0,0 +1,11 @@
+#!/bin/bash
+
+exp_dir=experiments/ClearSep_audioset_32k
+
+resume_ckpt="${exp_dir}/checkpoints/last.ckpt"
+
+python train.py $exp_dir \
+  --resume_ckpt $resume_ckpt \
+  --use_cuda \
+  --gpu_ids 0 1 2 3
+

train.py

@@ -0,0 +1,118 @@
+import os
+import logging
+import torch
+import torch.utils.data
+import pytorch_lightning as pl
+import laion_clap
+from pytorch_lightning.loggers import TensorBoardLogger
+from pytorch_lightning.callbacks import ModelCheckpoint
+from model.CLAPSep_decoder import HTSAT_Decoder
+from model.CLAPSep import LightningModule
+import argparse
+from helpers import utils as local_utils
+from dataset import CLAPSepDataSet, CLAPSepDataEngineDataSet
+
+import wandb
+from pytorch_lightning.loggers import WandbLogger
+
+
+def main(args):
+    torch.set_float32_matmul_precision('medium')
+    # Load dataset
+    data_train = CLAPSepDataEngineDataSet(**args.train_data)
+    # data_train = CLAPSepDataSet(**args.train_data)
+    logging.info("Loaded train dataset at %s containing %d elements" %
+                 (args.train_data['data_list'], len(data_train)))
+    data_val = CLAPSepDataSet(**args.val_data)
+    logging.info("Loaded test dataset at %s containing %d elements" %
+                 (args.val_data['data_list'], len(data_val)))
+    train_loader = torch.utils.data.DataLoader(data_train,
+                                               batch_size=args.batch_size,
+                                               shuffle=True,
+                                               num_workers=args.n_workers,
+                                               pin_memory=True)
+    val_loader = torch.utils.data.DataLoader(data_val,
+                                             batch_size=args.eval_batch_size,
+                                             shuffle=False,
+                                             num_workers=args.n_workers,
+                                             pin_memory=True)
+
+    clap_model = laion_clap.CLAP_Module(enable_fusion=False, amodel='HTSAT-base', device='cpu')
+    clap_model.load_ckpt(args.clap_path)
+    decoder = HTSAT_Decoder(**args.model)
+    lightning_module = LightningModule(clap_model, decoder, lr=args.optim['lr'],
+                                       use_lora=args.lora,
+                                       rank=args.lora_rank,
+                                       nfft=args.nfft,)
+    
+    checkpoint_callback = ModelCheckpoint(dirpath=os.path.join(args.exp_dir, 'checkpoints'),
+                                          filename="{epoch:02d}-{step}-{val_loss:.2f}",
+                                          monitor="val_loss",
+                                          mode="max",
+                                          save_top_k=3,
+                                          every_n_train_steps=args.save_ckpt_every_steps,
+                                          save_last=True)
+    logger = TensorBoardLogger(args.exp_dir)
+    # wandb_logger = WandbLogger(project='clapsep')
+    # wandb_logger = WandbLogger(project='clapsep', id='', resume='must')
+    # distributed_backend = "ddp_find_unused_parameters_true"
+    distributed_backend = "ddp"
+    trainer = pl.Trainer(
+        default_root_dir=args.exp_dir,
+        devices=args.gpu_ids if args.use_cuda else "auto",
+        accelerator="gpu" if args.use_cuda else "cpu",
+        benchmark=True,
+        gradient_clip_val=5.0,
+        precision='bf16-mixed',
+        limit_train_batches=1.0,
+        max_epochs=args.epochs,
+        strategy=distributed_backend,
+        logger=logger,
+        callbacks=[checkpoint_callback],
+    )
+
+    if os.path.exists(args.resume_ckpt):
+        print('Load resume ckpt:', args.resume_ckpt)
+        trainer.fit(model=lightning_module, train_dataloaders=train_loader, val_dataloaders=val_loader,
+                    ckpt_path=args.resume_ckpt)
+    elif os.path.exists(args.init_ckpt):
+        print('Load init ckpt:', args.init_ckpt)
+        weights = torch.load(args.init_ckpt, map_location='cpu')['state_dict']
+        lightning_module.load_state_dict(weights, strict=False)
+        trainer.fit(model=lightning_module, train_dataloaders=train_loader, val_dataloaders=val_loader)
+    else:
+        print('Training from scratch')
+        trainer.fit(model=lightning_module, train_dataloaders=train_loader, val_dataloaders=val_loader)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    # Data Params
+    parser.add_argument('exp_dir', type=str,
+                        default='./experiments/CLAPSep_base',
+                        help="Path to save checkpoints and logs.")
+    parser.add_argument('--init_ckpt', type=str, default='')
+    parser.add_argument('--resume_ckpt', type=str, default='')
+
+    parser.add_argument('--multi_label_training', dest='multi_label_training', action='store_true',
+                        help="Whether to multi label training")
+    
+    parser.add_argument('--use_cuda', dest='use_cuda', action='store_true',
+                        help="Whether to use cuda")
+    parser.add_argument('--gpu_ids', nargs='+', type=int, default=None,
+                        help="List of GPU ids used for training. "
+                             "Eg., --gpu_ids 2 4. All GPUs are used by default.")
+
+    args = parser.parse_args()
+
+    # Set the random seed for reproducible experiments
+    pl.seed_everything(114514)
+    # Set up checkpoints
+    if not os.path.exists(args.exp_dir):
+        os.makedirs(args.exp_dir)
+
+    # Load model and training params
+    params = local_utils.Params(os.path.join(args.exp_dir, 'config.json'))
+    for k, v in params.__dict__.items():
+        vars(args)[k] = v
+    main(args)