cleanup: remove dead code (training/, audio modality, chunk_inference, unused imports)

app.py

@@ -1,5 +1,5 @@
 import sys
-sys.stdout.reconfigure(line_buffering=True)
+sys.stdout.reconfigure(line_buffering=True)  # flush print() calls immediately so HF Space logs are live
 
 import os
 import tempfile
@@ -8,7 +8,6 @@ import traceback
 
 import gradio as gr
 import librosa
-import numpy as np
 import torch
 import soundfile as sf
 from huggingface_hub import hf_hub_download
@@ -17,12 +16,8 @@ from pyharp import ModelCard, build_endpoint
 from models.resunet import ResUNet30
 from models.clap_encoder import CLAP_Encoder
 
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-REPO_ID = "Audio-AGI/AudioSep"
-REPO_TYPE = "space"
-
 # Model loading state — populated by the background thread.
 ss_model = None
 query_encoder = None
@@ -33,24 +28,19 @@ model_error = None
 def load_model():
     global ss_model, query_encoder, model_loading, model_error
     try:
-        print(f"Downloading AudioSep checkpoint from {REPO_ID}...", flush=True)
+        print("Downloading checkpoints...")
         main_ckpt_path = hf_hub_download(
-            repo_id=REPO_ID,
-            repo_type=REPO_TYPE,
+            repo_id="Audio-AGI/AudioSep",
+            repo_type="space",
             filename="checkpoint/audiosep_base_4M_steps.ckpt",
         )
-
-        print(f"Downloading CLAP checkpoint from {REPO_ID}...", flush=True)
         clap_ckpt_path = hf_hub_download(
-            repo_id=REPO_ID,
-            repo_type=REPO_TYPE,
+            repo_id="Audio-AGI/AudioSep",
+            repo_type="space",
             filename="checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt",
         )
 
-        print("Loading CLAP encoder...", flush=True)
         query_encoder = CLAP_Encoder(pretrained_path=clap_ckpt_path).eval()
-
-        print("Loading ResUNet30 separation model...", flush=True)
         _ss_model = ResUNet30(input_channels=1, output_channels=1, condition_size=512)
 
         # Load weights from the Lightning checkpoint — keys are prefixed with "ss_model."
@@ -64,10 +54,10 @@ def load_model():
         _ss_model.eval().to(DEVICE)
 
         ss_model = _ss_model
-        print("Model loaded successfully.", flush=True)
+        print("Model loaded successfully.")
     except Exception as e:
         model_error = str(e)
-        print(f"Error loading model: {traceback.format_exc()}", flush=True)
+        print(f"Error loading model: {traceback.format_exc()}")
     finally:
         model_loading = False
 
@@ -84,15 +74,15 @@ model_card = ModelCard(
 
 
 @torch.inference_mode()
-def process_fn(audio_path: str, text_query: str) -> str:
+def process_fn(audio_path: str, text_query: str):
     if model_loading:
         raise gr.Error("Model is still loading, please wait a moment and try again.")
     if ss_model is None:
         raise gr.Error(f"Model failed to load: {model_error}")
 
-    print(f"Separating [{audio_path}] with query [{text_query}]", flush=True)
+    print(f"Separating [{audio_path}] with query [{text_query}]")
 
-    mixture, _ = librosa.load(audio_path, sr=32000, mono=True)
+    mixture, _ = librosa.load(audio_path, sr=32000, mono=True)  # model expects 32k audio
 
     conditions = query_encoder.get_query_embed(
         modality="text",

models/CLAP/training/audioset_textmap.npy

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

models/CLAP/training/data.py

@@ -1,966 +0,0 @@
-import ast
-import json
-import logging
-import math
-import os
-import random
-import io
-import copy
-import tempfile
-from dataclasses import dataclass
-from functools import partial
-from pathlib import Path
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchvision.transforms
-import braceexpand
-import soundfile as sf
-from PIL import Image
-from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler
-from torch.utils.data.distributed import DistributedSampler
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-try:
-    import torchaudio
-except ImportError:
-    torchaudio = None
-
-from models.CLAP.open_clip import tokenize
-
-
-def tokenizer(text):
-    return tokenize(text).squeeze(0)
-
-
-from transformers import RobertaTokenizer
-
-tokenize = RobertaTokenizer.from_pretrained("roberta-base")
-
-
-def tokenizer(text):
-    result = tokenize(
-        text,
-        padding="max_length",
-        truncation=True,
-        max_length=77,
-        return_tensors="pt",
-    )
-    return {k: v.squeeze(0) for k, v in result.items()}
-
-
-# initizlied the audioset map
-_AUDIOSET_MAP_PATH = os.path.join(Path(__file__).parent, "audioset_textmap.npy")
-_AUDIOSET_MAP = np.load(_AUDIOSET_MAP_PATH, allow_pickle=True)
-
-
-def int16_to_float32(x):
-    return (x / 32767.0).astype(np.float32)
-
-
-def float32_to_int16(x):
-    x = np.clip(x, a_min=-1.0, a_max=1.0)
-    return (x * 32767.0).astype(np.int16)
-
-
-# For Toy Dataset
-class ToyDataset(Dataset):
-    def __init__(self, index_path, ipc, config, eval_mode=False):
-        """Toy Dataset for testing the audioset input with text labels
-        Parameters
-        ----------
-            index_path: str
-                the link to the h5 file of each audio
-            idc: str
-                the link to the npy file, the number of samples in each class
-            config: dict
-                the audio cfg file
-           eval_model (bool): to indicate if the dataset is a testing dataset
-        """
-        self.audio_cfg = config["audio_cfg"]
-        self.text_cfg = config["text_cfg"]
-        self.fp = h5py.File(index_path, "r")
-        self.ipc = np.load(ipc, allow_pickle=True)
-        self.total_size = len(self.fp["audio_name"])
-        self.classes_num = self.audio_cfg["class_num"]
-        self.eval_mode = eval_mode
-
-        if not eval_mode:
-            self.generate_queue()
-        else:
-            self.queue = []
-            for i in range(self.total_size):
-                target = self.fp["target"][i]
-                if np.sum(target) > 0:
-                    self.queue.append(i)
-            self.total_size = len(self.queue)
-        logging.info("total dataset size: %d" % (self.total_size))
-        logging.info("class num: %d" % (self.classes_num))
-
-    def time_shifting(self, x):
-        frame_num = len(x)
-        shift_len = random.randint(0, frame_num - 1)
-        new_sample = np.concatenate([x[shift_len:], x[:shift_len]], axis=0)
-        return new_sample
-
-    def generate_queue(self):
-        self.queue = []
-        while len(self.queue) < self.total_size:
-            class_set = [*range(self.classes_num)]
-            random.shuffle(class_set)
-            self.queue += [
-                self.ipc[d][random.randint(0, len(self.ipc[d]) - 1)] for d in class_set
-            ]
-        self.queue = self.queue[: self.total_size]
-
-        logging.info("queue regenerated:%s" % (self.queue[-5:]))
-
-    def crop_wav(self, x):
-        crop_size = self.audio_cfg["crop_size"]
-        crop_pos = random.randint(0, len(x) - crop_size - 1)
-        return x[crop_pos : crop_pos + crop_size]
-
-    def prompt_text(self, target):
-        events = _AUDIOSET_MAP[np.where(target > 0)]
-        event_text = "The sounds of " + ", ".join(events[:-1]) + " and " + events[-1]
-        text = tokenize(event_text)[0]
-        return text
-
-    def __getitem__(self, index):
-        """Load waveform, text, and target of an audio clip
-
-        Parameters
-        ----------
-            index: int
-                the index number
-        Return
-        ------
-            output: dict {
-                "hdf5_path": str,
-                "index_in_hdf5": int,
-                "audio_name": str,
-                "waveform": list (audio_length,),
-                "target": list (class_num, ),
-                "text": torch.tensor (context_length,)
-            }
-                the output dictionary
-        """
-        s_index = self.queue[index]
-
-        audio_name = self.fp["audio_name"][s_index].decode()
-        # Hardcode here CHANGE
-        hdf5_path = (
-            self.fp["hdf5_path"][s_index]
-            .decode()
-            .replace(
-                "../workspace",
-                "/home/la/kechen/Research/ke_zsasp/workspace",
-            )
-        )
-        r_idx = self.fp["index_in_hdf5"][s_index]
-        target = self.fp["target"][s_index].astype(np.float32)
-        text = self.prompt_text(target)
-        with h5py.File(hdf5_path, "r") as f:
-            waveform = int16_to_float32(f["waveform"][r_idx])[
-                : self.audio_cfg["clip_samples"]
-            ]
-        assert (
-            len(waveform) == self.audio_cfg["clip_samples"]
-        ), "The sample length is not match"
-        # Time shift
-        # if (self.config.enable_time_shift) and (not self.eval_mode):
-        #     waveform = self.time_shifting(waveform)
-        # # Label Enhance
-        # if (self.config.crop_size is not None) and (not self.eval_mode):
-        #     waveform = self.crop_wav(waveform)
-        # # the label enhance rate is fixed 0.5
-        # if (self.config.enable_label_enhance) and (not self.eval_mode) and random.random() < 0.5:
-        #     kidx = np.where(target)[0]
-        #     for k in kidx:
-        #         for add_key in self.class_map[k][1]:
-        #             target[add_key] = 1.0
-        #         if len(self.class_map[k][2]) > 0:
-        #             add_key = random.choice(self.class_map[k][2])
-        #             target[add_key] = 1.0
-
-        # missing the text input
-        mel_spec = get_mel(torch.from_numpy(waveform), self.audio_cfg)[None, :, :]
-        mel_spec = (
-            torch.cat(
-                [mel_spec, mel_spec.clone(), mel_spec.clone(), mel_spec.clone()], dim=0
-            )
-            .cpu()
-            .numpy()
-        )
-        longer = random.choice([True, False])
-        if longer == False:
-            mel_spec[1:, :, :] = 0.0
-        data_dict = {
-            "hdf5_path": hdf5_path,
-            "index_in_hdf5": r_idx,
-            "audio_name": audio_name,
-            "waveform": waveform,
-            "class_label": target,
-            "text": text,
-            "longer": longer,
-            "mel_fusion": mel_spec,
-        }
-        return data_dict
-
-    def __len__(self):
-        return self.total_size
-
-
-class CsvDataset(Dataset):
-    def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t"):
-        logging.debug(f"Loading csv data from {input_filename}.")
-        df = pd.read_csv(input_filename, sep=sep)
-
-        self.images = df[img_key].tolist()
-        self.captions = df[caption_key].tolist()
-        self.transforms = transforms
-        logging.debug("Done loading data.")
-
-    def __len__(self):
-        return len(self.captions)
-
-    def __getitem__(self, idx):
-        images = self.transforms(Image.open(str(self.images[idx])))
-        texts = tokenize([str(self.captions[idx])])[0]
-        return images, texts
-
-
-@dataclass
-class DataInfo:
-    dataloader: DataLoader
-    sampler: DistributedSampler
-
-
-def preprocess_txt(text):
-    return tokenize([str(text)])[0]
-
-
-def get_dataset_size(shards, sizefilepath_=None, is_local=True):
-    if isinstance(shards, list):
-        size_list = []
-        for s in shards:
-            size_list.append(
-                get_dataset_size(s, sizefilepath_=sizefilepath_, is_local=is_local)[0]
-            )
-    else:
-        if not is_local:
-            for n in dataset_split.keys():
-                if n in shards.split("/"):
-                    break
-            for s in dataset_split[n]:
-                if s in shards.split("/"):
-                    break
-            sizefilepath_ = f"./json_files/{n}/{s}/sizes.json"
-        shards_list = list(braceexpand.braceexpand(shards))
-        dir_path = os.path.dirname(shards)
-        if sizefilepath_ is not None:
-            sizes = json.load(open(sizefilepath_, "r"))
-            total_size = sum(
-                [
-                    int(sizes[os.path.basename(shard.replace(".tar -", ".tar"))])
-                    for shard in shards_list
-                ]
-            )
-        else:
-            sizes_filename = os.path.join(dir_path, "sizes.json")
-            len_filename = os.path.join(dir_path, "__len__")
-            if os.path.exists(sizes_filename):
-                sizes = json.load(open(sizes_filename, "r"))
-                total_size = sum(
-                    [int(sizes[os.path.basename(shard)]) for shard in shards_list]
-                )
-            elif os.path.exists(len_filename):
-                # FIXME this used to be eval(open(...)) but that seemed rather unsafe
-                total_size = ast.literal_eval(open(len_filename, "r").read())
-            else:
-                raise Exception(
-                    "Cannot find sizes file for dataset. Please specify the path to the file."
-                )
-                # total_size = None  # num samples undefined
-                # some common dataset sizes (at time of authors last download)
-                # cc3m-train: 2905954
-                # cc12m: 10968539
-                # LAION-400m: 407332084
-        num_shards = len(shards_list)
-    if isinstance(shards, list):
-        return sum(size_list), len(shards)
-    else:
-        return total_size, num_shards
-
-
-def get_imagenet(args, preprocess_fns, split):
-    assert split in ["train", "val", "v2"]
-    is_train = split == "train"
-    preprocess_train, preprocess_val = preprocess_fns
-
-    if split == "v2":
-        from imagenetv2_pytorch import ImageNetV2Dataset
-
-        dataset = ImageNetV2Dataset(location=args.imagenet_v2, transform=preprocess_val)
-    else:
-        if is_train:
-            data_path = args.imagenet_train
-            preprocess_fn = preprocess_train
-        else:
-            data_path = args.imagenet_val
-            preprocess_fn = preprocess_val
-        assert data_path
-
-        dataset = datasets.ImageFolder(data_path, transform=preprocess_fn)
-
-    if is_train:
-        idxs = np.zeros(len(dataset.targets))
-        target_array = np.array(dataset.targets)
-        k = 50
-        for c in range(1000):
-            m = target_array == c
-            n = len(idxs[m])
-            arr = np.zeros(n)
-            arr[:k] = 1
-            np.random.shuffle(arr)
-            idxs[m] = arr
-
-        idxs = idxs.astype("int")
-        sampler = SubsetRandomSampler(np.where(idxs)[0])
-    else:
-        sampler = None
-
-    dataloader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        num_workers=args.workers,
-        sampler=sampler,
-    )
-
-    return DataInfo(dataloader, sampler)
-
-
-def count_samples(dataloader):
-    os.environ["WDS_EPOCH"] = "0"
-    n_elements, n_batches = 0, 0
-    for images, texts in dataloader:
-        n_batches += 1
-        n_elements += len(images)
-        assert len(images) == len(texts)
-    return n_elements, n_batches
-
-
-def filter_no_caption(sample):
-    return "txt" in sample
-
-
-def log_and_continue(exn):
-    """Call in an exception handler to ignore any exception, isssue a warning, and continue."""
-    logging.warning(f"Handling webdataset error ({repr(exn)}). Ignoring.")
-    return True
-
-
-_SHARD_SHUFFLE_SIZE = 2000
-_SHARD_SHUFFLE_INITIAL = 500
-_SAMPLE_SHUFFLE_SIZE = 5000
-_SAMPLE_SHUFFLE_INITIAL = 1000
-
-
-def sample_prop(sizefile, inputs, proportion, is_local=True):
-    """
-    Sample a proportion of the data.
-    """
-    file_path_dict = {
-        os.path.split(inputs[i])[1]: os.path.split(inputs[i])[0]
-        for i in range(len(inputs))
-    }
-    sampled_filepath_dict = {}
-    sampled_size_dict = {}
-    if not is_local:
-        if os.path.exists("sizes.json"):
-            os.remove("sizes.json")
-        wget.download(sizefile, "sizes.json")
-        sizefile = "sizes.json"
-    with open(sizefile, "r", encoding="UTF-8") as f:
-        load_dict = json.load(f)
-    L = int(len(file_path_dict) * proportion)
-    subkeys = random.sample(file_path_dict.keys(), L)
-    for k in subkeys:
-        sampled_size_dict[k] = load_dict[k]
-        sampled_filepath_dict[k] = file_path_dict[k]
-    return (
-        sum(sampled_size_dict.values()),
-        L,
-        [os.path.join(v, k) for k, v in sampled_filepath_dict.items()],
-        sampled_size_dict,
-    )
-
-
-def get_mel(audio_data, audio_cfg):
-    # mel shape: (n_mels, T)
-    mel = torchaudio.transforms.MelSpectrogram(
-        sample_rate=audio_cfg["sample_rate"],
-        n_fft=audio_cfg["window_size"],
-        win_length=audio_cfg["window_size"],
-        hop_length=audio_cfg["hop_size"],
-        center=True,
-        pad_mode="reflect",
-        power=2.0,
-        norm=None,
-        onesided=True,
-        n_mels=64,
-        f_min=audio_cfg["fmin"],
-        f_max=audio_cfg["fmax"],
-    ).to(audio_data.device)
-    mel = mel(audio_data)
-    # Align to librosa:
-    # librosa_melspec = librosa.feature.melspectrogram(
-    #     waveform,
-    #     sr=audio_cfg['sample_rate'],
-    #     n_fft=audio_cfg['window_size'],
-    #     hop_length=audio_cfg['hop_size'],
-    #     win_length=audio_cfg['window_size'],
-    #     center=True,
-    #     pad_mode="reflect",
-    #     power=2.0,
-    #     n_mels=64,
-    #     norm=None,
-    #     htk=True,
-    #     f_min=audio_cfg['fmin'],
-    #     f_max=audio_cfg['fmax']
-    # )
-    # we use log mel spectrogram as input
-    mel = torchaudio.transforms.AmplitudeToDB(top_db=None)(mel)
-    return mel.T  # (T, n_mels)
-
-
-def get_audio_features(
-    sample, audio_data, max_len, data_truncating, data_filling, audio_cfg
-):
-    """
-    Calculate and add audio features to sample.
-    Sample: a dict containing all the data of current sample.
-    audio_data: a tensor of shape (T) containing audio data.
-    max_len: the maximum length of audio data.
-    data_truncating: the method of truncating data.
-    data_filling: the method of filling data.
-    audio_cfg: a dict containing audio configuration. Comes from model_cfg['audio_cfg'].
-    """
-    with torch.no_grad():
-        if len(audio_data) > max_len:
-            if data_truncating == "rand_trunc":
-                longer = torch.tensor([True])
-            elif data_truncating == "fusion":
-                # fusion
-                mel = get_mel(audio_data, audio_cfg)
-                # split to three parts
-                chunk_frames = (
-                    max_len // audio_cfg["hop_size"] + 1
-                )  # the +1 related to how the spectrogram is computed
-                total_frames = mel.shape[0]
-                if chunk_frames == total_frames:
-                    # there is a corner case where the audio length is
-                    # larger than max_len but smaller than max_len+hop_size.
-                    # In this case, we just use the whole audio.
-                    mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
-                    sample["mel_fusion"] = mel_fusion
-                    longer = torch.tensor([False])
-                else:
-                    ranges = np.array_split(
-                        list(range(0, total_frames - chunk_frames + 1)), 3
-                    )
-                    # print('total_frames-chunk_frames:', total_frames-chunk_frames,
-                    #       'len(audio_data):', len(audio_data),
-                    #       'chunk_frames:', chunk_frames,
-                    #       'total_frames:', total_frames)
-                    if len(ranges[1]) == 0:
-                        # if the audio is too short, we just use the first chunk
-                        ranges[1] = [0]
-                    if len(ranges[2]) == 0:
-                        # if the audio is too short, we just use the first chunk
-                        ranges[2] = [0]
-                    # randomly choose index for each part
-                    idx_front = np.random.choice(ranges[0])
-                    idx_middle = np.random.choice(ranges[1])
-                    idx_back = np.random.choice(ranges[2])
-                    # select mel
-                    mel_chunk_front = mel[idx_front : idx_front + chunk_frames, :]
-                    mel_chunk_middle = mel[idx_middle : idx_middle + chunk_frames, :]
-                    mel_chunk_back = mel[idx_back : idx_back + chunk_frames, :]
-
-                    # shrink the mel
-                    mel_shrink = torchvision.transforms.Resize(size=[chunk_frames, 64])(
-                        mel[None]
-                    )[0]
-                    # logging.info(f"mel_shrink.shape: {mel_shrink.shape}")
-
-                    # stack
-                    mel_fusion = torch.stack(
-                        [mel_chunk_front, mel_chunk_middle, mel_chunk_back, mel_shrink],
-                        dim=0,
-                    )
-                    sample["mel_fusion"] = mel_fusion
-                    longer = torch.tensor([True])
-            else:
-                raise NotImplementedError(
-                    f"data_truncating {data_truncating} not implemented"
-                )
-            # random crop to max_len (for compatibility)
-            overflow = len(audio_data) - max_len
-            idx = np.random.randint(0, overflow + 1)
-            audio_data = audio_data[idx : idx + max_len]
-
-        else:  # padding if too short
-            if len(audio_data) < max_len:  # do nothing if equal
-                if data_filling == "repeatpad":
-                    n_repeat = int(max_len / len(audio_data))
-                    audio_data = audio_data.repeat(n_repeat)
-                    # audio_data = audio_data.unsqueeze(0).unsqueeze(0).unsqueeze(0)
-                    # audio_data = F.interpolate(audio_data,size=max_len,mode="bicubic")[0,0,0]
-                    audio_data = F.pad(
-                        audio_data,
-                        (0, max_len - len(audio_data)),
-                        mode="constant",
-                        value=0,
-                    )
-                elif data_filling == "pad":
-                    audio_data = F.pad(
-                        audio_data,
-                        (0, max_len - len(audio_data)),
-                        mode="constant",
-                        value=0,
-                    )
-                elif data_filling == "repeat":
-                    n_repeat = int(max_len / len(audio_data))
-                    audio_data = audio_data.repeat(n_repeat + 1)[:max_len]
-                else:
-                    raise NotImplementedError(
-                        f"data_filling {data_filling} not implemented"
-                    )
-            if data_truncating == "fusion":
-                mel = get_mel(audio_data, audio_cfg)
-                mel_fusion = torch.stack([mel, mel, mel, mel], dim=0)
-                sample["mel_fusion"] = mel_fusion
-            longer = torch.tensor([False])
-
-    sample["longer"] = longer
-    sample["waveform"] = audio_data
-
-    return sample
-
-
-def preprocess(
-    sample,
-    audio_ext,
-    text_ext,
-    max_len,
-    audio_cfg,
-    class_index_dict=None,
-    data_filling="pad",
-    data_truncating="rand_trunc",
-    text_augment_selection=None,
-):
-    """
-    Preprocess a single sample for wdsdataloader.
-    """
-    audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
-    audio_data = int16_to_float32(float32_to_int16(audio_data))
-    audio_data = torch.tensor(audio_data).float()
-
-    # TODO: (yusong) to be include in the future
-    # # if torchaudio not installed, use soundfile to load audio
-    # if torchaudio is None:
-    #     audio_data, orig_sr = sf.read(io.BytesIO(sample[audio_ext]))
-    #     audio_data = torch.tensor(audio_data).float()
-    # else:
-    #     # https://github.com/webdataset/webdataset/blob/main/webdataset/autodecode.py
-    #     with tempfile.TemporaryDirectory() as dirname:
-    #         os.makedirs(dirname, exist_ok=True)
-    #         fname = os.path.join(dirname, f"file.flac")
-    #         with open(fname, "wb") as stream:
-    #             stream.write(sample[audio_ext])
-    #         audio_data, orig_sr = torchaudio.load(fname)
-    #         audio_data = audio_data[0, :].float()
-
-    sample = get_audio_features(
-        sample, audio_data, max_len, data_truncating, data_filling, audio_cfg
-    )
-    del sample[audio_ext]
-
-    try:
-        json_dict_raw = json.loads(sample[text_ext].decode("utf-8"))
-    except:
-        print("sample[__url__]:", sample["__url__"])
-
-    # For selecting augmented text from dataset
-    if text_augment_selection is None or text_augment_selection == "none":
-        texts = json_dict_raw["text"]
-    elif text_augment_selection == "all":
-        if "text_augment_all" in json_dict_raw.keys():
-            texts = json_dict_raw["text_augment_all"]
-        else:
-            texts = json_dict_raw["text"]
-    elif text_augment_selection == "augment_only":
-        if "text_augment_all" in json_dict_raw.keys():
-            if json_dict_raw["text_augment_t5"] is None:
-                texts = json_dict_raw["text"]
-            else:
-                texts = json_dict_raw["text_augment_t5"]
-        else:
-            texts = json_dict_raw["text"]
-    else:
-        raise NotImplementedError(
-            f"text_augment_selection {text_augment_selection} not implemented"
-        )
-    sample["full_text"] = texts
-
-    if isinstance(texts, list) and isinstance(texts[0], str) and len(texts) > 1:
-        texts = random.choice(texts)
-    sample["raw_text"] = texts
-    sample["text"] = tokenizer(texts)  # text shape: [num_token]
-    if class_index_dict is not None:
-        # https://stackoverflow.com/questions/48004243/how-to-share-large-read-only-dictionary-list-across-processes-in-multiprocessing
-        # https://stackoverflow.com/questions/45693949/storing-strings-in-a-multiprocessing-sharedctypes-array
-        # key, val = class_index_dict
-        # key = key[:].split('\n')
-        # _dict = {k: v for k, v in zip(key, val)}
-        sample["class_label"] = np.zeros(len(class_index_dict.keys()))
-        for x in json_dict_raw["tag"]:
-            sample["class_label"][class_index_dict[x]] = 1
-        sample["class_label"] = torch.tensor(sample["class_label"]).float()
-    del sample[text_ext]
-    sample["audio_name"] = sample["__key__"].split("/")[-1] + "." + audio_ext
-    sample["text_name"] = sample["__key__"].split("/")[-1] + "." + text_ext
-    sample["audio_orig_sr"] = orig_sr
-    return sample
-
-
-def collate_fn(batch):
-    """
-    Collate function for wdsdataloader.
-    batch: a list of dict, each dict is a sample
-    """
-    # concatenate values in each dictionary. if it is a tensor, concatenate. if it is a list, extend.
-    batch_dict = {}
-    for k in batch[0].keys():
-        if isinstance(batch[0][k], dict):  # dealwith bert tokenizer output
-            batch_dict[k] = {}
-            for kk in batch[0][k].keys():
-                tmp = []
-                for i in range(len(batch)):
-                    tmp.append(batch[i][k][kk])
-                batch_dict[k][kk] = torch.vstack(tmp)
-        elif isinstance(batch[0][k], torch.Tensor):
-            batch_dict[k] = torch.stack([sample[k] for sample in batch])
-        elif isinstance(batch[0][k], np.ndarray):
-            batch_dict[k] = torch.tensor(np.stack([sample[k] for sample in batch]))
-        else:
-            batch_dict[k] = [sample[k] for sample in batch]
-    return batch_dict
-
-
-def get_wds_dataset(
-    args,
-    model_cfg,
-    is_train,
-    audio_ext="flac",
-    text_ext="json",
-    max_len=480000,
-    proportion=1.0,
-    sizefilepath_=None,
-    is_local=None,
-):
-    """
-    Get a dataset for wdsdataloader.
-    """
-    if is_local is None and (not args.remotedata is None):
-        is_local = not args.remotedata
-
-    input_shards = args.train_data if is_train else args.val_data
-    assert input_shards is not None
-
-    if not sizefilepath_ is None:
-        sizefilepath = sizefilepath_
-    else:
-        sizefilepath = os.path.join(os.path.dirname(input_shards[0]), "sizes.json")
-
-    if proportion != 1.0:
-        num_samples, num_shards, input_shards, _ = sample_prop(
-            sizefilepath, input_shards, proportion, is_local=is_local
-        )
-    else:
-        num_samples, num_shards = get_dataset_size(
-            input_shards, sizefilepath_=sizefilepath_, is_local=is_local
-        )
-
-    if not num_samples:
-        if is_train:
-            num_samples = args.train_num_samples
-            if not num_samples:
-                raise RuntimeError(
-                    "Currently, number of dataset samples must be specified for training dataset. "
-                    "Please specify via `--train-num-samples` if no dataset length info present."
-                )
-        else:
-            num_samples = (
-                args.val_num_samples or 0
-            )  # eval will just exhaust the iterator if not specified
-
-    pipeline = [wds.SimpleShardList(input_shards)]
-    # at this point we have an iterator over all the shards
-    # TODO: (yusong): add a if statement of distributed. If not, we don't need to split_by_node
-    if is_train or args.parallel_eval:
-        pipeline.extend(
-            [
-                wds.detshuffle(
-                    bufsize=_SHARD_SHUFFLE_SIZE,
-                    initial=_SHARD_SHUFFLE_INITIAL,
-                    seed=args.seed,
-                ),
-                wds.split_by_node,
-                wds.split_by_worker,
-                # at this point, we have an iterator over the shards assigned to each worker at each node
-                wds.tarfile_to_samples(handler=log_and_continue),
-                wds.shuffle(
-                    bufsize=_SAMPLE_SHUFFLE_SIZE,
-                    initial=_SAMPLE_SHUFFLE_INITIAL,
-                    rng=random.Random(args.seed),
-                ),
-                # wds.repeatedly,  # FIXME determine if this is beneficial
-            ]
-        )
-    else:
-        pipeline.extend(
-            [
-                wds.split_by_worker,
-                # at this point, we have an iterator over the shards assigned to each worker
-                wds.tarfile_to_samples(handler=log_and_continue),
-            ]
-        )
-    pipeline.append(
-        wds.map(
-            partial(
-                preprocess,
-                audio_ext=audio_ext,
-                text_ext=text_ext,
-                max_len=max_len,
-                audio_cfg=model_cfg["audio_cfg"],
-                class_index_dict=copy.deepcopy(args.class_index_dict),
-                data_filling=args.data_filling,
-                data_truncating=args.data_truncating,
-                text_augment_selection=args.text_augment_selection,
-            )
-        ),
-    )
-
-    pipeline.append(
-        wds.batched(
-            args.batch_size,
-            partial=not (is_train or args.parallel_eval),
-            collation_fn=collate_fn,
-        )
-    )
-
-    dataset = wds.DataPipeline(*pipeline)
-    if is_train or args.parallel_eval:
-        # (yusong): Currently parallel evaluation will be not precise as we are repeat the last few samples.
-        # (yusong): See comments below.
-        # roll over and repeat a few samples to get same number of full batches on each node
-        global_batch_size = args.batch_size * args.world_size
-        num_batches = math.ceil(num_samples / global_batch_size)
-        num_workers = max(1, args.workers)
-        num_worker_batches = math.ceil(
-            num_batches / num_workers
-        )  # per dataloader worker
-        num_batches = num_worker_batches * num_workers
-        num_samples = num_batches * global_batch_size
-        dataset = dataset.with_epoch(
-            num_worker_batches
-        )  # each worker is iterating over this
-    else:
-        # last batches are partial, eval is done on single (master) node
-        num_batches = math.ceil(num_samples / args.batch_size)
-
-    kwargs = {}
-    if args.horovod:  # multi-node training on summit
-        kwargs["multiprocessing_context"] = "forkserver"
-
-    dataloader = wds.WebLoader(
-        dataset, batch_size=None, shuffle=False, num_workers=args.workers, **kwargs
-    )
-
-    # FIXME not clear which approach is better, with_epoch before vs after dataloader?
-    # hoping to resolve via https://github.com/webdataset/webdataset/issues/169
-    # if is_train:
-    #     # roll over and repeat a few samples to get same number of full batches on each node
-    #     global_batch_size = args.batch_size * args.world_size
-    #     num_batches = math.ceil(num_samples / global_batch_size)
-    #     num_workers = max(1, args.workers)
-    #     num_batches = math.ceil(num_batches / num_workers) * num_workers
-    #     num_samples = num_batches * global_batch_size
-    #     dataloader = dataloader.with_epoch(num_batches)
-    # else:
-    #     # last batches are partial, eval is done on single (master) node
-    #     num_batches = math.ceil(num_samples / args.batch_size)
-
-    # add meta-data to dataloader instance for convenience
-    dataloader.num_batches = num_batches
-    dataloader.num_samples = num_samples
-
-    return DataInfo(dataloader, None)
-
-
-def wds_batch_list2dict(
-    batch,
-    keys=[
-        "__url__",
-        "__key__",
-        "waveform",
-        "text",
-        "raw_text",
-        "audio_name",
-        "text_name",
-        "audio_orig_sr",
-    ],
-):
-    """
-    Return a dictionary of the batch, with keys as the names of the fields.
-    """
-    assert len(keys) == len(
-        batch
-    ), "batch must have same number of keys as keys argument"
-    return {keys[i]: batch[i] for i in range(len(batch))}
-
-
-def get_csv_dataset(args, preprocess_fn, is_train):
-    input_filename = args.train_data if is_train else args.val_data
-    assert input_filename
-    dataset = CsvDataset(
-        input_filename,
-        preprocess_fn,
-        img_key=args.csv_img_key,
-        caption_key=args.csv_caption_key,
-        sep=args.csv_separator,
-    )
-    num_samples = len(dataset)
-    sampler = DistributedSampler(dataset) if args.distributed and is_train else None
-    shuffle = is_train and sampler is None
-
-    dataloader = DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        shuffle=shuffle,
-        num_workers=args.workers,
-        pin_memory=True,
-        sampler=sampler,
-        drop_last=is_train,
-    )
-    dataloader.num_samples = num_samples
-    dataloader.num_batches = len(dataloader)
-
-    return DataInfo(dataloader, sampler)
-
-
-def get_toy_dataset(args, model_cfg, is_train):
-    index_path = args.train_data if is_train else args.val_data
-    ipc_path = args.train_ipc if is_train else args.val_ipc
-    assert index_path and ipc_path
-    eval_mode = not is_train
-    dataset = ToyDataset(index_path, ipc_path, model_cfg, eval_mode=eval_mode)
-
-    num_samples = len(dataset)
-    sampler = (
-        DistributedSampler(dataset, shuffle=False)
-        if args.distributed and is_train
-        else None
-    )
-
-    dataloader = DataLoader(
-        dataset,
-        batch_size=args.batch_size,
-        shuffle=False,
-        num_workers=args.workers,
-        sampler=sampler,
-        drop_last=is_train,
-    )
-    dataloader.num_samples = num_samples
-    dataloader.num_batches = len(dataloader)
-
-    return DataInfo(dataloader, sampler)
-
-
-def get_dataset_fn(data_path, dataset_type):
-    if dataset_type == "webdataset":
-        return get_wds_dataset
-    elif dataset_type == "csv":
-        return get_csv_dataset
-    elif dataset_type == "auto":
-        ext = data_path.split(".")[-1]
-        if ext in ["csv", "tsv"]:
-            return get_csv_dataset
-        elif ext in ["tar"]:
-            return get_wds_dataset
-        else:
-            raise ValueError(
-                f"Tried to figure out dataset type, but failed for extention {ext}."
-            )
-    elif dataset_type == "toy":
-        return get_toy_dataset
-    else:
-        raise ValueError(f"Unsupported dataset type: {dataset_type}")
-
-
-def get_data(args, model_cfg):
-    data = {}
-
-    args.class_index_dict = load_class_label(args.class_label_path)
-
-    if args.datasetinfos is None:
-        args.datasetinfos = ["train", "unbalanced_train", "balanced_train"]
-    if args.dataset_type == "webdataset":
-        args.train_data = get_tar_path_from_dataset_name(
-            args.datasetnames,
-            args.datasetinfos,
-            islocal=not args.remotedata,
-            proportion=args.dataset_proportion,
-            dataset_path=args.datasetpath,
-            full_dataset=args.full_train_dataset,
-        )
-
-        if args.full_train_dataset is None:
-            args.full_train_dataset = []
-        if args.exclude_eval_dataset is None:
-            args.exclude_eval_dataset = []
-        excluded_eval_datasets = args.full_train_dataset + args.exclude_eval_dataset
-
-        val_dataset_names = (
-            [n for n in args.datasetnames if n not in excluded_eval_datasets]
-            if excluded_eval_datasets
-            else args.datasetnames
-        )
-        args.val_dataset_names = val_dataset_names
-        args.val_data = get_tar_path_from_dataset_name(
-            val_dataset_names,
-            ["valid", "test", "eval"],
-            islocal=not args.remotedata,
-            proportion=1,
-            dataset_path=args.datasetpath,
-            full_dataset=None,
-        )
-
-    if args.train_data:
-        data["train"] = get_dataset_fn(args.train_data, args.dataset_type)(
-            args, model_cfg, is_train=True
-        )
-
-    if args.val_data:
-        data["val"] = get_dataset_fn(args.val_data, args.dataset_type)(
-            args, model_cfg, is_train=False
-        )
-
-    return data

models/CLAP/training/distributed.py

@@ -1,150 +0,0 @@
-import os
-
-import torch
-import socket
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-
-def is_global_master(args):
-    return args.rank == 0
-
-
-def is_local_master(args):
-    return args.local_rank == 0
-
-
-def is_master(args, local=False):
-    return is_local_master(args) if local else is_global_master(args)
-
-
-def is_using_horovod():
-    # NOTE w/ horovod run, OMPI vars should be set, but w/ SLURM PMI vars will be set
-    # Differentiating between horovod and DDP use via SLURM may not be possible, so horovod arg still required...
-    ompi_vars = ["OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"]
-    pmi_vars = ["PMI_RANK", "PMI_SIZE"]
-    if all([var in os.environ for var in ompi_vars]) or all(
-        [var in os.environ for var in pmi_vars]
-    ):
-        return True
-    else:
-        return False
-
-
-def is_using_distributed():
-    if "WORLD_SIZE" in os.environ:
-        return int(os.environ["WORLD_SIZE"]) > 1
-    if "SLURM_NTASKS" in os.environ:
-        return int(os.environ["SLURM_NTASKS"]) > 1
-    return False
-
-
-def world_info_from_env():
-    local_rank = 0
-    for v in (
-        "SLURM_LOCALID",
-        "MPI_LOCALRANKID",
-        "OMPI_COMM_WORLD_LOCAL_RANK",
-        "LOCAL_RANK",
-    ):
-        if v in os.environ:
-            local_rank = int(os.environ[v])
-            break
-    global_rank = 0
-    for v in ("SLURM_PROCID", "PMI_RANK", "OMPI_COMM_WORLD_RANK", "RANK"):
-        if v in os.environ:
-            global_rank = int(os.environ[v])
-            break
-    world_size = 1
-    for v in ("SLURM_NTASKS", "PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "WORLD_SIZE"):
-        if v in os.environ:
-            world_size = int(os.environ[v])
-            break
-
-    return local_rank, global_rank, world_size
-
-
-def init_distributed_device(args):
-    # Distributed training = training on more than one GPU.
-    # Works in both single and multi-node scenarios.
-    args.distributed = False
-    args.world_size = 1
-    args.rank = 0  # global rank
-    args.local_rank = 0
-    if args.horovod:
-        assert hvd is not None, "Horovod is not installed"
-        hvd.init()
-        world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
-        world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
-        local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
-        args.local_rank = local_rank
-        args.rank = world_rank
-        args.world_size = world_size
-        # args.local_rank = int(hvd.local_rank())
-        # args.rank = hvd.rank()
-        # args.world_size = hvd.size()
-        args.distributed = True
-        os.environ["LOCAL_RANK"] = str(args.local_rank)
-        os.environ["RANK"] = str(args.rank)
-        os.environ["WORLD_SIZE"] = str(args.world_size)
-        print(
-            f"Distributed training: local_rank={args.local_rank}, "
-            f"rank={args.rank}, world_size={args.world_size}, "
-            f"hostname={socket.gethostname()}, pid={os.getpid()}"
-        )
-    elif is_using_distributed():
-        if "SLURM_PROCID" in os.environ:
-            # DDP via SLURM
-            args.local_rank, args.rank, args.world_size = world_info_from_env()
-            # SLURM var -> torch.distributed vars in case needed
-            os.environ["LOCAL_RANK"] = str(args.local_rank)
-            os.environ["RANK"] = str(args.rank)
-            os.environ["WORLD_SIZE"] = str(args.world_size)
-            torch.distributed.init_process_group(
-                backend=args.dist_backend,
-                init_method=args.dist_url,
-                world_size=args.world_size,
-                rank=args.rank,
-            )
-        elif "OMPI_COMM_WORLD_SIZE" in os.environ:  # using Summit cluster
-            world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"])
-            world_rank = int(os.environ["OMPI_COMM_WORLD_RANK"])
-            local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])
-            args.local_rank = local_rank
-            args.rank = world_rank
-            args.world_size = world_size
-            torch.distributed.init_process_group(
-                backend=args.dist_backend,
-                init_method=args.dist_url,
-                world_size=args.world_size,
-                rank=args.rank,
-            )
-        else:
-            # DDP via torchrun, torch.distributed.launch
-            args.local_rank, _, _ = world_info_from_env()
-            torch.distributed.init_process_group(
-                backend=args.dist_backend, init_method=args.dist_url
-            )
-            args.world_size = torch.distributed.get_world_size()
-            args.rank = torch.distributed.get_rank()
-        args.distributed = True
-        print(
-            f"Distributed training: local_rank={args.local_rank}, "
-            f"rank={args.rank}, world_size={args.world_size}, "
-            f"hostname={socket.gethostname()}, pid={os.getpid()}"
-        )
-
-    if torch.cuda.is_available():
-        if args.distributed and not args.no_set_device_rank:
-            device = "cuda:%d" % args.local_rank
-        else:
-            device = "cuda:0"
-        torch.cuda.set_device(device)
-    else:
-        device = "cpu"
-    args.device = device
-    device = torch.device(device)
-    return device

models/CLAP/training/imagenet_zeroshot_data.py

@@ -1,1088 +0,0 @@
-# NOTE: This script is currently not supported for CLAP.
-
-imagenet_classnames = [
-    "tench",
-    "goldfish",
-    "great white shark",
-    "tiger shark",
-    "hammerhead shark",
-    "electric ray",
-    "stingray",
-    "rooster",
-    "hen",
-    "ostrich",
-    "brambling",
-    "goldfinch",
-    "house finch",
-    "junco",
-    "indigo bunting",
-    "American robin",
-    "bulbul",
-    "jay",
-    "magpie",
-    "chickadee",
-    "American dipper",
-    "kite (bird of prey)",
-    "bald eagle",
-    "vulture",
-    "great grey owl",
-    "fire salamander",
-    "smooth newt",
-    "newt",
-    "spotted salamander",
-    "axolotl",
-    "American bullfrog",
-    "tree frog",
-    "tailed frog",
-    "loggerhead sea turtle",
-    "leatherback sea turtle",
-    "mud turtle",
-    "terrapin",
-    "box turtle",
-    "banded gecko",
-    "green iguana",
-    "Carolina anole",
-    "desert grassland whiptail lizard",
-    "agama",
-    "frilled-necked lizard",
-    "alligator lizard",
-    "Gila monster",
-    "European green lizard",
-    "chameleon",
-    "Komodo dragon",
-    "Nile crocodile",
-    "American alligator",
-    "triceratops",
-    "worm snake",
-    "ring-necked snake",
-    "eastern hog-nosed snake",
-    "smooth green snake",
-    "kingsnake",
-    "garter snake",
-    "water snake",
-    "vine snake",
-    "night snake",
-    "boa constrictor",
-    "African rock python",
-    "Indian cobra",
-    "green mamba",
-    "sea snake",
-    "Saharan horned viper",
-    "eastern diamondback rattlesnake",
-    "sidewinder rattlesnake",
-    "trilobite",
-    "harvestman",
-    "scorpion",
-    "yellow garden spider",
-    "barn spider",
-    "European garden spider",
-    "southern black widow",
-    "tarantula",
-    "wolf spider",
-    "tick",
-    "centipede",
-    "black grouse",
-    "ptarmigan",
-    "ruffed grouse",
-    "prairie grouse",
-    "peafowl",
-    "quail",
-    "partridge",
-    "african grey parrot",
-    "macaw",
-    "sulphur-crested cockatoo",
-    "lorikeet",
-    "coucal",
-    "bee eater",
-    "hornbill",
-    "hummingbird",
-    "jacamar",
-    "toucan",
-    "duck",
-    "red-breasted merganser",
-    "goose",
-    "black swan",
-    "tusker",
-    "echidna",
-    "platypus",
-    "wallaby",
-    "koala",
-    "wombat",
-    "jellyfish",
-    "sea anemone",
-    "brain coral",
-    "flatworm",
-    "nematode",
-    "conch",
-    "snail",
-    "slug",
-    "sea slug",
-    "chiton",
-    "chambered nautilus",
-    "Dungeness crab",
-    "rock crab",
-    "fiddler crab",
-    "red king crab",
-    "American lobster",
-    "spiny lobster",
-    "crayfish",
-    "hermit crab",
-    "isopod",
-    "white stork",
-    "black stork",
-    "spoonbill",
-    "flamingo",
-    "little blue heron",
-    "great egret",
-    "bittern bird",
-    "crane bird",
-    "limpkin",
-    "common gallinule",
-    "American coot",
-    "bustard",
-    "ruddy turnstone",
-    "dunlin",
-    "common redshank",
-    "dowitcher",
-    "oystercatcher",
-    "pelican",
-    "king penguin",
-    "albatross",
-    "grey whale",
-    "killer whale",
-    "dugong",
-    "sea lion",
-    "Chihuahua",
-    "Japanese Chin",
-    "Maltese",
-    "Pekingese",
-    "Shih Tzu",
-    "King Charles Spaniel",
-    "Papillon",
-    "toy terrier",
-    "Rhodesian Ridgeback",
-    "Afghan Hound",
-    "Basset Hound",
-    "Beagle",
-    "Bloodhound",
-    "Bluetick Coonhound",
-    "Black and Tan Coonhound",
-    "Treeing Walker Coonhound",
-    "English foxhound",
-    "Redbone Coonhound",
-    "borzoi",
-    "Irish Wolfhound",
-    "Italian Greyhound",
-    "Whippet",
-    "Ibizan Hound",
-    "Norwegian Elkhound",
-    "Otterhound",
-    "Saluki",
-    "Scottish Deerhound",
-    "Weimaraner",
-    "Staffordshire Bull Terrier",
-    "American Staffordshire Terrier",
-    "Bedlington Terrier",
-    "Border Terrier",
-    "Kerry Blue Terrier",
-    "Irish Terrier",
-    "Norfolk Terrier",
-    "Norwich Terrier",
-    "Yorkshire Terrier",
-    "Wire Fox Terrier",
-    "Lakeland Terrier",
-    "Sealyham Terrier",
-    "Airedale Terrier",
-    "Cairn Terrier",
-    "Australian Terrier",
-    "Dandie Dinmont Terrier",
-    "Boston Terrier",
-    "Miniature Schnauzer",
-    "Giant Schnauzer",
-    "Standard Schnauzer",
-    "Scottish Terrier",
-    "Tibetan Terrier",
-    "Australian Silky Terrier",
-    "Soft-coated Wheaten Terrier",
-    "West Highland White Terrier",
-    "Lhasa Apso",
-    "Flat-Coated Retriever",
-    "Curly-coated Retriever",
-    "Golden Retriever",
-    "Labrador Retriever",
-    "Chesapeake Bay Retriever",
-    "German Shorthaired Pointer",
-    "Vizsla",
-    "English Setter",
-    "Irish Setter",
-    "Gordon Setter",
-    "Brittany dog",
-    "Clumber Spaniel",
-    "English Springer Spaniel",
-    "Welsh Springer Spaniel",
-    "Cocker Spaniel",
-    "Sussex Spaniel",
-    "Irish Water Spaniel",
-    "Kuvasz",
-    "Schipperke",
-    "Groenendael dog",
-    "Malinois",
-    "Briard",
-    "Australian Kelpie",
-    "Komondor",
-    "Old English Sheepdog",
-    "Shetland Sheepdog",
-    "collie",
-    "Border Collie",
-    "Bouvier des Flandres dog",
-    "Rottweiler",
-    "German Shepherd Dog",
-    "Dobermann",
-    "Miniature Pinscher",
-    "Greater Swiss Mountain Dog",
-    "Bernese Mountain Dog",
-    "Appenzeller Sennenhund",
-    "Entlebucher Sennenhund",
-    "Boxer",
-    "Bullmastiff",
-    "Tibetan Mastiff",
-    "French Bulldog",
-    "Great Dane",
-    "St. Bernard",
-    "husky",
-    "Alaskan Malamute",
-    "Siberian Husky",
-    "Dalmatian",
-    "Affenpinscher",
-    "Basenji",
-    "pug",
-    "Leonberger",
-    "Newfoundland dog",
-    "Great Pyrenees dog",
-    "Samoyed",
-    "Pomeranian",
-    "Chow Chow",
-    "Keeshond",
-    "brussels griffon",
-    "Pembroke Welsh Corgi",
-    "Cardigan Welsh Corgi",
-    "Toy Poodle",
-    "Miniature Poodle",
-    "Standard Poodle",
-    "Mexican hairless dog (xoloitzcuintli)",
-    "grey wolf",
-    "Alaskan tundra wolf",
-    "red wolf or maned wolf",
-    "coyote",
-    "dingo",
-    "dhole",
-    "African wild dog",
-    "hyena",
-    "red fox",
-    "kit fox",
-    "Arctic fox",
-    "grey fox",
-    "tabby cat",
-    "tiger cat",
-    "Persian cat",
-    "Siamese cat",
-    "Egyptian Mau",
-    "cougar",
-    "lynx",
-    "leopard",
-    "snow leopard",
-    "jaguar",
-    "lion",
-    "tiger",
-    "cheetah",
-    "brown bear",
-    "American black bear",
-    "polar bear",
-    "sloth bear",
-    "mongoose",
-    "meerkat",
-    "tiger beetle",
-    "ladybug",
-    "ground beetle",
-    "longhorn beetle",
-    "leaf beetle",
-    "dung beetle",
-    "rhinoceros beetle",
-    "weevil",
-    "fly",
-    "bee",
-    "ant",
-    "grasshopper",
-    "cricket insect",
-    "stick insect",
-    "cockroach",
-    "praying mantis",
-    "cicada",
-    "leafhopper",
-    "lacewing",
-    "dragonfly",
-    "damselfly",
-    "red admiral butterfly",
-    "ringlet butterfly",
-    "monarch butterfly",
-    "small white butterfly",
-    "sulphur butterfly",
-    "gossamer-winged butterfly",
-    "starfish",
-    "sea urchin",
-    "sea cucumber",
-    "cottontail rabbit",
-    "hare",
-    "Angora rabbit",
-    "hamster",
-    "porcupine",
-    "fox squirrel",
-    "marmot",
-    "beaver",
-    "guinea pig",
-    "common sorrel horse",
-    "zebra",
-    "pig",
-    "wild boar",
-    "warthog",
-    "hippopotamus",
-    "ox",
-    "water buffalo",
-    "bison",
-    "ram (adult male sheep)",
-    "bighorn sheep",
-    "Alpine ibex",
-    "hartebeest",
-    "impala (antelope)",
-    "gazelle",
-    "arabian camel",
-    "llama",
-    "weasel",
-    "mink",
-    "European polecat",
-    "black-footed ferret",
-    "otter",
-    "skunk",
-    "badger",
-    "armadillo",
-    "three-toed sloth",
-    "orangutan",
-    "gorilla",
-    "chimpanzee",
-    "gibbon",
-    "siamang",
-    "guenon",
-    "patas monkey",
-    "baboon",
-    "macaque",
-    "langur",
-    "black-and-white colobus",
-    "proboscis monkey",
-    "marmoset",
-    "white-headed capuchin",
-    "howler monkey",
-    "titi monkey",
-    "Geoffroy's spider monkey",
-    "common squirrel monkey",
-    "ring-tailed lemur",
-    "indri",
-    "Asian elephant",
-    "African bush elephant",
-    "red panda",
-    "giant panda",
-    "snoek fish",
-    "eel",
-    "silver salmon",
-    "rock beauty fish",
-    "clownfish",
-    "sturgeon",
-    "gar fish",
-    "lionfish",
-    "pufferfish",
-    "abacus",
-    "abaya",
-    "academic gown",
-    "accordion",
-    "acoustic guitar",
-    "aircraft carrier",
-    "airliner",
-    "airship",
-    "altar",
-    "ambulance",
-    "amphibious vehicle",
-    "analog clock",
-    "apiary",
-    "apron",
-    "trash can",
-    "assault rifle",
-    "backpack",
-    "bakery",
-    "balance beam",
-    "balloon",
-    "ballpoint pen",
-    "Band-Aid",
-    "banjo",
-    "baluster / handrail",
-    "barbell",
-    "barber chair",
-    "barbershop",
-    "barn",
-    "barometer",
-    "barrel",
-    "wheelbarrow",
-    "baseball",
-    "basketball",
-    "bassinet",
-    "bassoon",
-    "swimming cap",
-    "bath towel",
-    "bathtub",
-    "station wagon",
-    "lighthouse",
-    "beaker",
-    "military hat (bearskin or shako)",
-    "beer bottle",
-    "beer glass",
-    "bell tower",
-    "baby bib",
-    "tandem bicycle",
-    "bikini",
-    "ring binder",
-    "binoculars",
-    "birdhouse",
-    "boathouse",
-    "bobsleigh",
-    "bolo tie",
-    "poke bonnet",
-    "bookcase",
-    "bookstore",
-    "bottle cap",
-    "hunting bow",
-    "bow tie",
-    "brass memorial plaque",
-    "bra",
-    "breakwater",
-    "breastplate",
-    "broom",
-    "bucket",
-    "buckle",
-    "bulletproof vest",
-    "high-speed train",
-    "butcher shop",
-    "taxicab",
-    "cauldron",
-    "candle",
-    "cannon",
-    "canoe",
-    "can opener",
-    "cardigan",
-    "car mirror",
-    "carousel",
-    "tool kit",
-    "cardboard box / carton",
-    "car wheel",
-    "automated teller machine",
-    "cassette",
-    "cassette player",
-    "castle",
-    "catamaran",
-    "CD player",
-    "cello",
-    "mobile phone",
-    "chain",
-    "chain-link fence",
-    "chain mail",
-    "chainsaw",
-    "storage chest",
-    "chiffonier",
-    "bell or wind chime",
-    "china cabinet",
-    "Christmas stocking",
-    "church",
-    "movie theater",
-    "cleaver",
-    "cliff dwelling",
-    "cloak",
-    "clogs",
-    "cocktail shaker",
-    "coffee mug",
-    "coffeemaker",
-    "spiral or coil",
-    "combination lock",
-    "computer keyboard",
-    "candy store",
-    "container ship",
-    "convertible",
-    "corkscrew",
-    "cornet",
-    "cowboy boot",
-    "cowboy hat",
-    "cradle",
-    "construction crane",
-    "crash helmet",
-    "crate",
-    "infant bed",
-    "Crock Pot",
-    "croquet ball",
-    "crutch",
-    "cuirass",
-    "dam",
-    "desk",
-    "desktop computer",
-    "rotary dial telephone",
-    "diaper",
-    "digital clock",
-    "digital watch",
-    "dining table",
-    "dishcloth",
-    "dishwasher",
-    "disc brake",
-    "dock",
-    "dog sled",
-    "dome",
-    "doormat",
-    "drilling rig",
-    "drum",
-    "drumstick",
-    "dumbbell",
-    "Dutch oven",
-    "electric fan",
-    "electric guitar",
-    "electric locomotive",
-    "entertainment center",
-    "envelope",
-    "espresso machine",
-    "face powder",
-    "feather boa",
-    "filing cabinet",
-    "fireboat",
-    "fire truck",
-    "fire screen",
-    "flagpole",
-    "flute",
-    "folding chair",
-    "football helmet",
-    "forklift",
-    "fountain",
-    "fountain pen",
-    "four-poster bed",
-    "freight car",
-    "French horn",
-    "frying pan",
-    "fur coat",
-    "garbage truck",
-    "gas mask or respirator",
-    "gas pump",
-    "goblet",
-    "go-kart",
-    "golf ball",
-    "golf cart",
-    "gondola",
-    "gong",
-    "gown",
-    "grand piano",
-    "greenhouse",
-    "radiator grille",
-    "grocery store",
-    "guillotine",
-    "hair clip",
-    "hair spray",
-    "half-track",
-    "hammer",
-    "hamper",
-    "hair dryer",
-    "hand-held computer",
-    "handkerchief",
-    "hard disk drive",
-    "harmonica",
-    "harp",
-    "combine harvester",
-    "hatchet",
-    "holster",
-    "home theater",
-    "honeycomb",
-    "hook",
-    "hoop skirt",
-    "gymnastic horizontal bar",
-    "horse-drawn vehicle",
-    "hourglass",
-    "iPod",
-    "clothes iron",
-    "carved pumpkin",
-    "jeans",
-    "jeep",
-    "T-shirt",
-    "jigsaw puzzle",
-    "rickshaw",
-    "joystick",
-    "kimono",
-    "knee pad",
-    "knot",
-    "lab coat",
-    "ladle",
-    "lampshade",
-    "laptop computer",
-    "lawn mower",
-    "lens cap",
-    "letter opener",
-    "library",
-    "lifeboat",
-    "lighter",
-    "limousine",
-    "ocean liner",
-    "lipstick",
-    "slip-on shoe",
-    "lotion",
-    "music speaker",
-    "loupe magnifying glass",
-    "sawmill",
-    "magnetic compass",
-    "messenger bag",
-    "mailbox",
-    "tights",
-    "one-piece bathing suit",
-    "manhole cover",
-    "maraca",
-    "marimba",
-    "mask",
-    "matchstick",
-    "maypole",
-    "maze",
-    "measuring cup",
-    "medicine cabinet",
-    "megalith",
-    "microphone",
-    "microwave oven",
-    "military uniform",
-    "milk can",
-    "minibus",
-    "miniskirt",
-    "minivan",
-    "missile",
-    "mitten",
-    "mixing bowl",
-    "mobile home",
-    "ford model t",
-    "modem",
-    "monastery",
-    "monitor",
-    "moped",
-    "mortar and pestle",
-    "graduation cap",
-    "mosque",
-    "mosquito net",
-    "vespa",
-    "mountain bike",
-    "tent",
-    "computer mouse",
-    "mousetrap",
-    "moving van",
-    "muzzle",
-    "metal nail",
-    "neck brace",
-    "necklace",
-    "baby pacifier",
-    "notebook computer",
-    "obelisk",
-    "oboe",
-    "ocarina",
-    "odometer",
-    "oil filter",
-    "pipe organ",
-    "oscilloscope",
-    "overskirt",
-    "bullock cart",
-    "oxygen mask",
-    "product packet / packaging",
-    "paddle",
-    "paddle wheel",
-    "padlock",
-    "paintbrush",
-    "pajamas",
-    "palace",
-    "pan flute",
-    "paper towel",
-    "parachute",
-    "parallel bars",
-    "park bench",
-    "parking meter",
-    "railroad car",
-    "patio",
-    "payphone",
-    "pedestal",
-    "pencil case",
-    "pencil sharpener",
-    "perfume",
-    "Petri dish",
-    "photocopier",
-    "plectrum",
-    "Pickelhaube",
-    "picket fence",
-    "pickup truck",
-    "pier",
-    "piggy bank",
-    "pill bottle",
-    "pillow",
-    "ping-pong ball",
-    "pinwheel",
-    "pirate ship",
-    "drink pitcher",
-    "block plane",
-    "planetarium",
-    "plastic bag",
-    "plate rack",
-    "farm plow",
-    "plunger",
-    "Polaroid camera",
-    "pole",
-    "police van",
-    "poncho",
-    "pool table",
-    "soda bottle",
-    "plant pot",
-    "potter's wheel",
-    "power drill",
-    "prayer rug",
-    "printer",
-    "prison",
-    "missile",
-    "projector",
-    "hockey puck",
-    "punching bag",
-    "purse",
-    "quill",
-    "quilt",
-    "race car",
-    "racket",
-    "radiator",
-    "radio",
-    "radio telescope",
-    "rain barrel",
-    "recreational vehicle",
-    "fishing casting reel",
-    "reflex camera",
-    "refrigerator",
-    "remote control",
-    "restaurant",
-    "revolver",
-    "rifle",
-    "rocking chair",
-    "rotisserie",
-    "eraser",
-    "rugby ball",
-    "ruler measuring stick",
-    "sneaker",
-    "safe",
-    "safety pin",
-    "salt shaker",
-    "sandal",
-    "sarong",
-    "saxophone",
-    "scabbard",
-    "weighing scale",
-    "school bus",
-    "schooner",
-    "scoreboard",
-    "CRT monitor",
-    "screw",
-    "screwdriver",
-    "seat belt",
-    "sewing machine",
-    "shield",
-    "shoe store",
-    "shoji screen / room divider",
-    "shopping basket",
-    "shopping cart",
-    "shovel",
-    "shower cap",
-    "shower curtain",
-    "ski",
-    "balaclava ski mask",
-    "sleeping bag",
-    "slide rule",
-    "sliding door",
-    "slot machine",
-    "snorkel",
-    "snowmobile",
-    "snowplow",
-    "soap dispenser",
-    "soccer ball",
-    "sock",
-    "solar thermal collector",
-    "sombrero",
-    "soup bowl",
-    "keyboard space bar",
-    "space heater",
-    "space shuttle",
-    "spatula",
-    "motorboat",
-    "spider web",
-    "spindle",
-    "sports car",
-    "spotlight",
-    "stage",
-    "steam locomotive",
-    "through arch bridge",
-    "steel drum",
-    "stethoscope",
-    "scarf",
-    "stone wall",
-    "stopwatch",
-    "stove",
-    "strainer",
-    "tram",
-    "stretcher",
-    "couch",
-    "stupa",
-    "submarine",
-    "suit",
-    "sundial",
-    "sunglasses",
-    "sunglasses",
-    "sunscreen",
-    "suspension bridge",
-    "mop",
-    "sweatshirt",
-    "swim trunks / shorts",
-    "swing",
-    "electrical switch",
-    "syringe",
-    "table lamp",
-    "tank",
-    "tape player",
-    "teapot",
-    "teddy bear",
-    "television",
-    "tennis ball",
-    "thatched roof",
-    "front curtain",
-    "thimble",
-    "threshing machine",
-    "throne",
-    "tile roof",
-    "toaster",
-    "tobacco shop",
-    "toilet seat",
-    "torch",
-    "totem pole",
-    "tow truck",
-    "toy store",
-    "tractor",
-    "semi-trailer truck",
-    "tray",
-    "trench coat",
-    "tricycle",
-    "trimaran",
-    "tripod",
-    "triumphal arch",
-    "trolleybus",
-    "trombone",
-    "hot tub",
-    "turnstile",
-    "typewriter keyboard",
-    "umbrella",
-    "unicycle",
-    "upright piano",
-    "vacuum cleaner",
-    "vase",
-    "vaulted or arched ceiling",
-    "velvet fabric",
-    "vending machine",
-    "vestment",
-    "viaduct",
-    "violin",
-    "volleyball",
-    "waffle iron",
-    "wall clock",
-    "wallet",
-    "wardrobe",
-    "military aircraft",
-    "sink",
-    "washing machine",
-    "water bottle",
-    "water jug",
-    "water tower",
-    "whiskey jug",
-    "whistle",
-    "hair wig",
-    "window screen",
-    "window shade",
-    "Windsor tie",
-    "wine bottle",
-    "airplane wing",
-    "wok",
-    "wooden spoon",
-    "wool",
-    "split-rail fence",
-    "shipwreck",
-    "sailboat",
-    "yurt",
-    "website",
-    "comic book",
-    "crossword",
-    "traffic or street sign",
-    "traffic light",
-    "dust jacket",
-    "menu",
-    "plate",
-    "guacamole",
-    "consomme",
-    "hot pot",
-    "trifle",
-    "ice cream",
-    "popsicle",
-    "baguette",
-    "bagel",
-    "pretzel",
-    "cheeseburger",
-    "hot dog",
-    "mashed potatoes",
-    "cabbage",
-    "broccoli",
-    "cauliflower",
-    "zucchini",
-    "spaghetti squash",
-    "acorn squash",
-    "butternut squash",
-    "cucumber",
-    "artichoke",
-    "bell pepper",
-    "cardoon",
-    "mushroom",
-    "Granny Smith apple",
-    "strawberry",
-    "orange",
-    "lemon",
-    "fig",
-    "pineapple",
-    "banana",
-    "jackfruit",
-    "cherimoya (custard apple)",
-    "pomegranate",
-    "hay",
-    "carbonara",
-    "chocolate syrup",
-    "dough",
-    "meatloaf",
-    "pizza",
-    "pot pie",
-    "burrito",
-    "red wine",
-    "espresso",
-    "tea cup",
-    "eggnog",
-    "mountain",
-    "bubble",
-    "cliff",
-    "coral reef",
-    "geyser",
-    "lakeshore",
-    "promontory",
-    "sandbar",
-    "beach",
-    "valley",
-    "volcano",
-    "baseball player",
-    "bridegroom",
-    "scuba diver",
-    "rapeseed",
-    "daisy",
-    "yellow lady's slipper",
-    "corn",
-    "acorn",
-    "rose hip",
-    "horse chestnut seed",
-    "coral fungus",
-    "agaric",
-    "gyromitra",
-    "stinkhorn mushroom",
-    "earth star fungus",
-    "hen of the woods mushroom",
-    "bolete",
-    "corn cob",
-    "toilet paper",
-]
-
-
-openai_imagenet_template = [
-    lambda c: f"a bad photo of a {c}.",
-    lambda c: f"a photo of many {c}.",
-    lambda c: f"a sculpture of a {c}.",
-    lambda c: f"a photo of the hard to see {c}.",
-    lambda c: f"a low resolution photo of the {c}.",
-    lambda c: f"a rendering of a {c}.",
-    lambda c: f"graffiti of a {c}.",
-    lambda c: f"a bad photo of the {c}.",
-    lambda c: f"a cropped photo of the {c}.",
-    lambda c: f"a tattoo of a {c}.",
-    lambda c: f"the embroidered {c}.",
-    lambda c: f"a photo of a hard to see {c}.",
-    lambda c: f"a bright photo of a {c}.",
-    lambda c: f"a photo of a clean {c}.",
-    lambda c: f"a photo of a dirty {c}.",
-    lambda c: f"a dark photo of the {c}.",
-    lambda c: f"a drawing of a {c}.",
-    lambda c: f"a photo of my {c}.",
-    lambda c: f"the plastic {c}.",
-    lambda c: f"a photo of the cool {c}.",
-    lambda c: f"a close-up photo of a {c}.",
-    lambda c: f"a black and white photo of the {c}.",
-    lambda c: f"a painting of the {c}.",
-    lambda c: f"a painting of a {c}.",
-    lambda c: f"a pixelated photo of the {c}.",
-    lambda c: f"a sculpture of the {c}.",
-    lambda c: f"a bright photo of the {c}.",
-    lambda c: f"a cropped photo of a {c}.",
-    lambda c: f"a plastic {c}.",
-    lambda c: f"a photo of the dirty {c}.",
-    lambda c: f"a jpeg corrupted photo of a {c}.",
-    lambda c: f"a blurry photo of the {c}.",
-    lambda c: f"a photo of the {c}.",
-    lambda c: f"a good photo of the {c}.",
-    lambda c: f"a rendering of the {c}.",
-    lambda c: f"a {c} in a video game.",
-    lambda c: f"a photo of one {c}.",
-    lambda c: f"a doodle of a {c}.",
-    lambda c: f"a close-up photo of the {c}.",
-    lambda c: f"a photo of a {c}.",
-    lambda c: f"the origami {c}.",
-    lambda c: f"the {c} in a video game.",
-    lambda c: f"a sketch of a {c}.",
-    lambda c: f"a doodle of the {c}.",
-    lambda c: f"a origami {c}.",
-    lambda c: f"a low resolution photo of a {c}.",
-    lambda c: f"the toy {c}.",
-    lambda c: f"a rendition of the {c}.",
-    lambda c: f"a photo of the clean {c}.",
-    lambda c: f"a photo of a large {c}.",
-    lambda c: f"a rendition of a {c}.",
-    lambda c: f"a photo of a nice {c}.",
-    lambda c: f"a photo of a weird {c}.",
-    lambda c: f"a blurry photo of a {c}.",
-    lambda c: f"a cartoon {c}.",
-    lambda c: f"art of a {c}.",
-    lambda c: f"a sketch of the {c}.",
-    lambda c: f"a embroidered {c}.",
-    lambda c: f"a pixelated photo of a {c}.",
-    lambda c: f"itap of the {c}.",
-    lambda c: f"a jpeg corrupted photo of the {c}.",
-    lambda c: f"a good photo of a {c}.",
-    lambda c: f"a plushie {c}.",
-    lambda c: f"a photo of the nice {c}.",
-    lambda c: f"a photo of the small {c}.",
-    lambda c: f"a photo of the weird {c}.",
-    lambda c: f"the cartoon {c}.",
-    lambda c: f"art of the {c}.",
-    lambda c: f"a drawing of the {c}.",
-    lambda c: f"a photo of the large {c}.",
-    lambda c: f"a black and white photo of a {c}.",
-    lambda c: f"the plushie {c}.",
-    lambda c: f"a dark photo of a {c}.",
-    lambda c: f"itap of a {c}.",
-    lambda c: f"graffiti of the {c}.",
-    lambda c: f"a toy {c}.",
-    lambda c: f"itap of my {c}.",
-    lambda c: f"a photo of a cool {c}.",
-    lambda c: f"a photo of a small {c}.",
-    lambda c: f"a tattoo of the {c}.",
-]

models/CLAP/training/infer_demo.py

@@ -1,109 +0,0 @@
-import sys
-
-sys.path.append(
-    "/mnt/fast/nobackup/users/hl01486/projects/contrastive_pretraining/CLAP/src"
-)
-
-import os
-import torch
-import librosa
-from open_clip import create_model
-from training.data import get_audio_features
-from training.data import int16_to_float32, float32_to_int16
-from transformers import RobertaTokenizer
-
-tokenize = RobertaTokenizer.from_pretrained("roberta-base")
-
-
-def tokenizer(text):
-    result = tokenize(
-        text,
-        padding="max_length",
-        truncation=True,
-        max_length=77,
-        return_tensors="pt",
-    )
-    return {k: v.squeeze(0) for k, v in result.items()}
-
-
-PRETRAINED_PATH = "/mnt/fast/nobackup/users/hl01486/projects/contrastive_pretraining/CLAP/assets/checkpoints/epoch_top_0_audioset_no_fusion.pt"
-WAVE_48k_PATH = "/mnt/fast/nobackup/users/hl01486/projects/contrastive_pretraining/CLAP/assets/audio/machine.wav"
-
-
-def infer_text():
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    precision = "fp32"
-    amodel = "HTSAT-tiny"  # or 'PANN-14'
-    tmodel = "roberta"  # the best text encoder in our training
-    enable_fusion = False  # False if you do not want to use the fusion model
-    fusion_type = "aff_2d"
-    pretrained = PRETRAINED_PATH
-
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision=precision,
-        device=device,
-        enable_fusion=enable_fusion,
-        fusion_type=fusion_type,
-    )
-    # load the text, can be a list (i.e. batch size)
-    text_data = ["I love the contrastive learning", "I love the pretrain model"]
-    # tokenize for roberta, if you want to tokenize for another text encoder, please refer to data.py#L43-90
-    text_data = tokenizer(text_data)
-
-    text_embed = model.get_text_embedding(text_data)
-    print(text_embed.size())
-
-
-def infer_audio():
-
-    device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    precision = "fp32"
-    amodel = "HTSAT-tiny"  # or 'PANN-14'
-    tmodel = "roberta"  # the best text encoder in our training
-    enable_fusion = False  # False if you do not want to use the fusion model
-    fusion_type = "aff_2d"
-    pretrained = PRETRAINED_PATH
-
-    model, model_cfg = create_model(
-        amodel,
-        tmodel,
-        pretrained,
-        precision=precision,
-        device=device,
-        enable_fusion=enable_fusion,
-        fusion_type=fusion_type,
-    )
-
-    # load the waveform of the shape (T,), should resample to 48000
-    audio_waveform, sr = librosa.load(WAVE_48k_PATH, sr=48000)
-    # quantize
-    audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
-    audio_waveform = torch.from_numpy(audio_waveform).float()
-    audio_dict = {}
-
-    # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
-    import ipdb
-
-    ipdb.set_trace()
-    audio_dict = get_audio_features(
-        audio_dict,
-        audio_waveform,
-        480000,
-        data_truncating="fusion",
-        data_filling="repeatpad",
-        audio_cfg=model_cfg["audio_cfg"],
-    )
-    # can send a list to the model, to process many audio tracks in one time (i.e. batch size)
-    audio_embed = model.get_audio_embedding([audio_dict])
-    print(audio_embed.size())
-    import ipdb
-
-    ipdb.set_trace()
-
-
-if __name__ == "__main__":
-    infer_text()
-    infer_audio()

models/CLAP/training/logger.py

@@ -1,30 +0,0 @@
-import logging
-
-
-def setup_logging(log_file, level, include_host=False):
-    if include_host:
-        import socket
-
-        hostname = socket.gethostname()
-        formatter = logging.Formatter(
-            f"%(asctime)s |  {hostname} | %(levelname)s | %(message)s",
-            datefmt="%Y-%m-%d,%H:%M:%S",
-        )
-    else:
-        formatter = logging.Formatter(
-            "%(asctime)s | %(levelname)s | %(message)s", datefmt="%Y-%m-%d,%H:%M:%S"
-        )
-
-    logging.root.setLevel(level)
-    loggers = [logging.getLogger(name) for name in logging.root.manager.loggerDict]
-    for logger in loggers:
-        logger.setLevel(level)
-
-    stream_handler = logging.StreamHandler()
-    stream_handler.setFormatter(formatter)
-    logging.root.addHandler(stream_handler)
-
-    if log_file:
-        file_handler = logging.FileHandler(filename=log_file)
-        file_handler.setFormatter(formatter)
-        logging.root.addHandler(file_handler)

models/CLAP/training/lp_main.py

@@ -1,670 +0,0 @@
-from cmath import cos
-from inspect import getargs
-import logging
-import os
-import random
-from datetime import datetime
-import bisect
-import copy
-from sched import scheduler
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-from torch import optim
-from torch.cuda.amp import GradScaler
-import faulthandler
-import pathlib
-import argparse
-import time
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-try:
-    import torch.utils.tensorboard as tensorboard
-except ImportError:
-    tensorboard = None
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-from open_clip import create_model_and_transforms, trace_model, create_model
-from training.data import get_data
-from training.params import parse_args
-from training.distributed import is_master, init_distributed_device, world_info_from_env
-from training.logger import setup_logging
-from training.scheduler import cosine_lr
-from training.lp_train import train_one_epoch, evaluate
-from open_clip.utils import get_tar_path_from_dataset_name, dataset_split, get_optimizer
-from open_clip.utils import load_p, load_class_label
-from open_clip.linear_probe import LinearProbe
-
-
-def maintain_ckpts(args, startidx, all_idx_len):
-    for i in reversed(range(startidx, all_idx_len)):
-        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
-            os.rename(
-                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
-            )
-    if os.path.exists(
-        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
-    ):
-        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
-    return
-
-
-def update_top_k_performance(
-    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True
-):
-    """
-    Record the top-k performance of the current epoch.
-    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
-    """
-    if isinstance(new_metrics_inputs, (list, tuple)):
-        new_metrics_inputs = np.mean(new_metrics_inputs)
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, dict):
-        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, (float, int)):
-        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
-        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
-        sorted_values = sorted(
-            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
-        )
-        sorted_values_ = copy.deepcopy(sorted_values)
-        sorted_values.append(new_metrics_inputs)
-        sorted_values = sorted(sorted_values, reverse=bignumbetter)
-        sorted_values = sorted_values[:-1]
-
-        if sorted_values == sorted_values_:
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-        else:
-            for i in range(len(sorted_keys)):
-                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
-                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
-                    update_flag[sorted_keys[i]] = True
-            for i in range(len(update_flag)):
-                if update_flag[i]:
-                    maintain_ckpts(args, i, len(sorted_keys))
-                    torch.save(
-                        ckpt,
-                        os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                    )
-                    break
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-
-
-# def updateifNone(a, b):
-#     a = b if None else a
-#     return a
-
-
-def is_pretrained_params(n):
-    return (
-        n.startswith("clap_model.transformer")
-        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
-        or n.startswith("clap_model.token_embedding")
-        or n.startswith("clap_model.ln_final")
-        or n.startswith("clap_model.logit_scale_t")
-    )
-
-
-def random_seed(seed=42, rank=0):
-    torch.manual_seed(seed + rank)
-    np.random.seed(seed + rank)
-    random.seed(seed + rank)
-
-
-def config_lp_optimizer(model, data, args):
-    # set wd-related params to 0 if use adam optimizer
-    if args.optimizer == "adam":
-        args.wd = 0
-        args.wd_pretrained = 0
-        args.wd_new = 0
-
-    in_clap = lambda n, p: n.startswith("clap_model")
-
-    named_parameters = list(model.named_parameters())
-
-    optimizer = {}
-    scheduler = {}
-
-    # freeze text encoder
-    text_freeze_parameters = [
-        p
-        for n, p in named_parameters
-        if n.startswith("clap_model.transformer")
-        or n in ["clap_model.positional_embedding", "clap_model.text_projection"]
-        or n.startswith("clap_model.token_embedding")
-        or n.startswith("clap_model.ln_final")
-    ]
-
-    if args.freeze_text:
-        logging.info("Freeze Text!!!!")
-        for k in text_freeze_parameters:
-            k.requires_grad = False
-
-    if not args.lp_freeze:
-        exclude = (
-            lambda n, p: p.ndim < 2
-            or "bn" in n
-            or "ln" in n
-            or "bias" in n
-            or "logit_scale" in n
-        )
-        include = lambda n, p: not exclude(n, p)
-
-        # (yusong): we do not split the learning rate anymore
-        # p for n, p in named_parameters if in_clap(n,p) and exclude(n, p) and p.requires_grad
-        gain_or_bias_params = [
-            p for n, p in named_parameters if exclude(n, p) and p.requires_grad
-        ]
-        # rest_params = [p for n, p in named_parameters if in_clap(n,p) and include(n, p) and p.requires_grad]
-        rest_params = [
-            p for n, p in named_parameters if include(n, p) and p.requires_grad
-        ]
-
-        if args.train_data is None:
-            optimizer = None
-            scheduler = None
-        else:
-            total_steps = data["train"].dataloader.num_batches * args.epochs
-
-            if args.split_opt:
-                for x in ["lr", "beta1", "beta2", "eps", "wd"]:
-                    for y in ["_new", "_pretrained"]:
-                        if getattr(args, x + y) is None:
-                            setattr(args, x + y, getattr(args, x))
-
-                gain_or_bias_pretrained_params = [
-                    p
-                    for n, p in named_parameters
-                    if (exclude(n, p) and p.requires_grad) and is_pretrained_params(n)
-                ]
-                rest_pretrained_params = [
-                    p
-                    for n, p in named_parameters
-                    if (include(n, p) and p.requires_grad) and is_pretrained_params(n)
-                ]
-                gain_or_bias_new_params = [
-                    p
-                    for n, p in named_parameters
-                    if (exclude(n, p) and p.requires_grad)
-                    and (not is_pretrained_params(n))
-                ]
-                rest_new_params = [
-                    p
-                    for n, p in named_parameters
-                    if (include(n, p) and p.requires_grad)
-                    and (not is_pretrained_params(n))
-                ]
-
-                pretrained_params_optimizer = get_optimizer(
-                    [
-                        {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
-                        {
-                            "params": rest_pretrained_params,
-                            "weight_decay": args.wd_pretrained,
-                        },
-                    ],
-                    lr=args.lr_pretrained,
-                    betas=(args.beta1_pretrained, args.beta2_pretrained),
-                    eps=args.eps_pretrained,
-                    momentum=args.momentum_pretrained,
-                    optimizer_name=args.optimizer,
-                )
-                pretrained_params_scheduler = cosine_lr(
-                    pretrained_params_optimizer,
-                    args.lr_pretrained,
-                    args.warmup,
-                    total_steps,
-                )
-
-                new_params_optimizer = get_optimizer(
-                    [
-                        {"params": gain_or_bias_new_params, "weight_decay": 0.0},
-                        {"params": rest_new_params, "weight_decay": args.wd_new},
-                    ],
-                    lr=args.lr_new,
-                    betas=(args.beta1_new, args.beta2_new),
-                    eps=args.eps_new,
-                    momentum=args.momentum_new,
-                    optimizer_name=args.optimizer,
-                )
-                new_params_scheduler = cosine_lr(
-                    new_params_optimizer, args.lr_new, args.warmup, total_steps
-                )
-
-                optimizer["text"] = pretrained_params_optimizer
-                optimizer["audio"] = new_params_optimizer
-                scheduler["text"] = pretrained_params_scheduler
-                scheduler["audio"] = new_params_scheduler
-
-                if args.horovod:
-                    pretrained_params_optimizer = hvd.DistributedOptimizer(
-                        pretrained_params_optimizer,
-                        named_parameters=model.named_parameters(),
-                    )
-                    new_params_optimizer = hvd.DistributedOptimizer(
-                        new_params_optimizer, named_parameters=model.named_parameters()
-                    )
-                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                    hvd.broadcast_optimizer_state(
-                        pretrained_params_optimizer, root_rank=0
-                    )
-                    hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
-            else:
-
-                optimizer["clap"] = get_optimizer(
-                    [
-                        {"params": gain_or_bias_params, "weight_decay": 0.0},
-                        {"params": rest_params, "weight_decay": args.wd},
-                    ],
-                    lr=args.lr,
-                    betas=(args.beta1, args.beta2),
-                    eps=args.eps,
-                    momentum=args.momentum,
-                    optimizer_name=args.optimizer,
-                )
-                scheduler["clap"] = cosine_lr(
-                    optimizer["clap"], args.lr, args.warmup, total_steps
-                )
-
-                if args.horovod:
-                    optimizer["clap"] = hvd.DistributedOptimizer(
-                        optimizer["clap"], named_parameters=model.named_parameters()
-                    )
-                    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                    hvd.broadcast_optimizer_state(optimizer["clap"], root_rank=0)
-
-    # linear probe optimizer
-    else:
-        lp_params = [
-            p for n, p in named_parameters if (not in_clap(n, p)) and p.requires_grad
-        ]
-        lp_optim = get_optimizer(
-            lp_params,
-            lr=args.lp_lr,
-            betas=(args.beta1, args.beta2),
-            eps=args.eps,
-            momentum=0.9,
-            optimizer_name=args.optimizer,
-        )
-        optimizer["lp"] = lp_optim
-
-    return optimizer, scheduler, text_freeze_parameters
-
-
-def main():
-    args = parse_args()
-
-    time.sleep(args.sleep)
-
-    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.amodel = args.amodel.replace("/", "-")
-    # download sizes.json file
-
-    # (yusong): the below two lines are for debug
-    # print("setting up faulthandler")
-    # faulthandler.register(10)
-
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-    args.class_index_dict = load_class_label(args.class_label_path)
-
-    # get the name of the experiments
-    if args.name is None:
-        args.name = "-".join(
-            [
-                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"linear_probe" f"model_{args.amodel}",
-                f"lr_{args.lr}",
-                f"b_{args.batch_size}",
-                f"j_{args.workers}",
-                f"p_{args.precision}",
-            ]
-        )
-
-    # discover initial world args early so we can log properly
-    args.distributed = False
-    args.local_rank, args.rank, args.world_size = world_info_from_env()
-
-    if args.remotedata and is_master(args):
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    args.log_path = None
-    if is_master(args, local=args.log_local):
-        log_base_path = os.path.join(args.logs, args.name)
-        os.makedirs(log_base_path, exist_ok=True)
-        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-        args.log_path = os.path.join(log_base_path, log_filename)
-
-        # avoid log dir in same name:
-        postfix = 0
-        while os.path.exists(args.log_path):
-            postfix += 1
-            log_base_path_new = log_base_path + "-" + str(postfix)
-            os.makedirs(log_base_path_new, exist_ok=True)
-            log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-            args.log_path = os.path.join(log_base_path_new, log_filename)
-            # print(
-            #     "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            # )
-            # return -1
-
-    # Set logger
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    setup_logging(args.log_path, args.log_level)
-
-    # fully initialize distributed device environment
-    device = init_distributed_device(args)
-
-    args.wandb = "wandb" in args.report_to or "all" in args.report_to
-    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
-    if is_master(args):
-        args.tensorboard_path = (
-            os.path.join(args.logs, args.name, "tensorboard")
-            if args.tensorboard
-            else ""
-        )
-        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
-        for dirname in [args.tensorboard_path, args.checkpoint_path]:
-            if dirname:
-                os.makedirs(dirname, exist_ok=True)
-    else:
-        args.tensorboard_path = ""
-        args.checkpoint_path = ""
-
-    if args.copy_codebase:
-        copy_codebase(args)
-
-    assert args.precision in ["amp", "fp16", "fp32"]
-    if args.precision == "fp16":
-        logging.warning(
-            "It is recommended to use AMP mixed-precision instead of FP16. "
-            "FP16 support needs further verification and tuning, especially for train."
-        )
-
-    if args.horovod:
-        logging.info(
-            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    elif args.distributed:
-        logging.info(
-            f"Running in distributed mode with multiple processes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    else:
-        logging.info(f"Running with a single process. Device {args.device}.")
-
-    logging.info(f"openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}")
-
-    # Create CLAP model
-    clap_model, clap_model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=False,
-        pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type,
-    )
-
-    args.lp_out_ch = len(list(args.class_index_dict.keys()))
-    # Linear Probe
-    logging.info(f"linear probe using mlp: {args.lp_mlp}")
-    logging.info(f"linear probe using freeze: {args.lp_freeze}")
-    logging.info(f"linear probe act layer: {args.lp_act}")
-    logging.info(f"linear probe out ch: {args.lp_out_ch}")
-    logging.info(f"linear probe learning rate (if applicable): {args.lp_lr}")
-    logging.info(f"linear probe loss func: {args.lp_loss}")
-    logging.info(f"linear probe lp_metrics: {args.lp_metrics}")
-
-    model = LinearProbe(
-        clap_model,
-        mlp=args.lp_mlp,
-        freeze=args.lp_freeze,
-        in_ch=512,
-        out_ch=args.lp_out_ch,
-        act=args.lp_act,
-    )  # in_ch is fixed (i.e., 512)
-    model = model.to(device)
-
-    if args.horovod:
-        with torch.no_grad():
-            for param in model.parameters():
-                param.set_(param.contiguous())
-
-    if args.trace:
-        model = trace_model(model, batch_size=args.batch_size, device=device)
-
-    if is_master(args):
-        logging.info("Linear Probe CLAP Model:")
-        logging.info(f"{str(clap_model)}")
-        logging.info("Params:")
-        params_file = os.path.join(args.logs, args.name, "params.txt")
-        with open(params_file, "w") as f:
-            for name in sorted(vars(args)):
-                val = getattr(args, name)
-                logging.info(f"  {name}: {val}")
-                f.write(f"{name}: {val}\n")
-
-    if args.distributed and not args.horovod:
-        if args.use_bn_sync:
-            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        ddp_args = {}
-        if args.ddp_static_graph:
-            # this doesn't exist in older PyTorch, arg only added if enabled
-            ddp_args["static_graph"] = True
-        model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[device], find_unused_parameters=True, **ddp_args
-        )
-
-    data = get_data(args, clap_model_cfg)
-    assert len(data), "At least one train or eval dataset must be specified."
-    if args.trace:
-        assert "train" not in data, "Cannot train with traced model"
-
-    optimizer, scheduler, text_freeze_parameters = config_lp_optimizer(
-        model, data, args
-    )
-
-    scaler = GradScaler() if args.precision == "amp" else None
-
-    # optionally resume from a checkpoint
-    start_epoch = 0
-    if args.resume is not None:
-        if os.path.isfile(args.resume):
-            checkpoint = torch.load(args.resume, map_location=device)
-            if "epoch" in checkpoint:
-                # resuming a train checkpoint w/ epoch and optimizer state
-                start_epoch = checkpoint["epoch"]
-                sd = checkpoint["state_dict"]
-                if not args.distributed and next(iter(sd.items()))[0].startswith(
-                    "module"
-                ):
-                    sd = {k[len("module.") :]: v for k, v in sd.items()}
-                model.load_state_dict(sd)
-                if args.split_opt:
-                    if optimizer is not None:
-                        for k, o_ in optimizer.items():
-                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
-                if optimizer is not None:
-                    optimizer.load_state_dict(checkpoint["optimizer"])
-                if scaler is not None and "scaler" in checkpoint:
-                    scaler.load_state_dict(checkpoint["scaler"])
-                logging.info(
-                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            else:
-                # loading a bare (model only) checkpoint for fine-tune or evaluation
-                model.load_state_dict(checkpoint)
-                logging.info(
-                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            if args.freeze_text:
-                print("Freeze Text!!!!")
-                for k in text_freeze_parameters:
-                    k.requires_grad = False
-        else:
-            logging.info("=> no checkpoint found at '{}'".format(args.resume))
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
-    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
-    writer = None
-    if args.save_logs and args.tensorboard:
-        assert tensorboard is not None, "Please install tensorboard."
-        writer = tensorboard.SummaryWriter(args.tensorboard_path)
-
-    if args.wandb and is_master(args):
-        assert wandb is not None, "Please install wandb."
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        wandb.init(
-            project="clap",
-            notes=args.wandb_notes,
-            name=args.wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
-        if args.debug:
-            wandb.watch(model, log="all")
-        wandb.save(params_file)
-        logging.debug("Finished loading wandb.")
-
-    if "train" not in data:
-        evaluate(model, data, start_epoch, args, writer)
-        return
-    elif start_epoch == 0 and "val" in data and not args.no_eval:
-        evaluate(model, data, 0, args, writer)
-    if args.save_top_performance:
-        current_top_k_ckpt_metrics = {
-            i: 0 for i in range(args.save_top_performance)
-        }  # initialize the top-k metric for ckpts to 0
-
-    for epoch in range(start_epoch, args.epochs):
-        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
-        if epoch == args.freeze_text_after:
-            print("Text pretrained parameters are freezed since this epoch.")
-            for k in text_freeze_parameters:
-                k.requires_grad = False
-        if is_master(args):
-            logging.info(f"Start epoch {epoch}")
-
-        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
-        completed_epoch = epoch + 1
-
-        if (
-            any(v in data for v in ("val", "imagenet-val", "imagenet-v2"))
-            and not args.no_eval
-        ):
-            metrics = evaluate(model, data, completed_epoch, args, writer)
-            if args.save_top_performance:
-                top_k_dataset = args.top_k_checkpoint_select_dataset
-                top_k_metric = args.top_k_checkpoint_select_metric
-                filtered_metrics = [
-                    v
-                    for k, v in metrics.items()
-                    if top_k_metric in k and top_k_dataset in k
-                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
-        # Saving checkpoints.
-        if args.save_logs:
-            opt_dict = {
-                k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
-            }
-            checkpoint_dict = {
-                "epoch": completed_epoch,
-                "name": args.name,
-                "state_dict": model.state_dict(),
-            }
-            checkpoint_dict.update(opt_dict)
-            if scaler is not None:
-                checkpoint_dict["scaler"] = scaler.state_dict()
-
-            if completed_epoch == args.epochs or (
-                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
-            ):
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_{completed_epoch}.pt"),
-                )
-            if args.save_most_recent:
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
-                )
-            if args.save_top_performance and not args.no_eval:
-                update_top_k_performance(
-                    filtered_metrics,
-                    current_top_k_ckpt_metrics,
-                    args,
-                    checkpoint_dict,
-                    bignumbetter=True,
-                )
-
-    if args.wandb and is_master(args):
-        wandb.finish()
-
-
-def copy_codebase(args):
-    from shutil import copytree, ignore_patterns
-
-    new_code_path = os.path.join(args.logs, args.name, "code")
-    if os.path.exists(new_code_path):
-        print(
-            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
-        )
-        return -1
-    print(f"Copying codebase to {new_code_path}")
-    current_code_path = os.path.realpath(__file__)
-    for _ in range(3):
-        current_code_path = os.path.dirname(current_code_path)
-    copytree(
-        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
-    )
-    print("Done copying code.")
-    return 1
-
-
-if __name__ == "__main__":
-    main()

models/CLAP/training/lp_train.py

@@ -1,301 +0,0 @@
-import json
-import logging
-import math
-import os
-import time
-from contextlib import suppress
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-from open_clip import LPLoss, LPMetrics, lp_gather_features
-from open_clip.utils import do_mixup, get_mix_lambda
-from .distributed import is_master
-from .zero_shot import zero_shot_eval
-
-
-class AverageMeter(object):
-    """Computes and stores the average and current value"""
-
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.val = 0
-        self.avg = 0
-        self.sum = 0
-        self.count = 0
-
-    def update(self, val, n=1):
-        self.val = val
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
-
-
-def unwrap_model(model):
-    if hasattr(model, "module"):
-        return model.module
-    else:
-        return model
-
-
-def train_one_epoch(
-    model,
-    data,
-    epoch,
-    optimizer,
-    scaler,
-    scheduler,
-    args,
-    tb_writer=None,
-    extra_suffix="",
-):
-    device = torch.device(args.device)
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    model.train()
-    loss = LPLoss(args.lp_loss)
-
-    dataloader, sampler = data["train"].dataloader, data["train"].sampler
-    if args.distributed and sampler is not None:
-        sampler.set_epoch(epoch)
-    num_batches_per_epoch = dataloader.num_batches
-    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
-
-    # for toy dataset
-    if args.dataset_type == "toy":
-        dataloader.dataset.generate_queue()
-
-    loss_m = AverageMeter()
-    batch_time_m = AverageMeter()
-    data_time_m = AverageMeter()
-    end = time.time()
-
-    for i, batch in enumerate(dataloader):
-        step = num_batches_per_epoch * epoch + i
-
-        if isinstance(scheduler, dict):
-            for s in scheduler.values():
-                s(step)
-        else:
-            scheduler(step)
-
-        audio = batch  # contains mel_spec, wavform, and longer list
-        class_label = batch["class_label"]
-        # audio = audio.to(device=device, non_blocking=True)
-        class_label = class_label.to(device=device, non_blocking=True)
-
-        if args.mixup:
-            # https://github.com/RetroCirce/HTS-Audio-Transformer/blob/main/utils.py#L146
-            mix_lambda = torch.from_numpy(
-                get_mix_lambda(0.5, len(audio["waveform"]))
-            ).to(device)
-            class_label = do_mixup(class_label, mix_lambda)
-        else:
-            mix_lambda = None
-
-        data_time_m.update(time.time() - end)
-        if isinstance(optimizer, dict):
-            for o_ in optimizer.values():
-                o_.zero_grad()
-        else:
-            optimizer.zero_grad()
-
-        with autocast():
-            pred = model(audio, mix_lambda=mix_lambda, device=device)
-            total_loss = loss(pred, class_label)
-
-        if isinstance(optimizer, dict):
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                for o_ in optimizer.values():
-                    if args.horovod:
-                        o_.synchronize()
-                        scaler.unscale_(o_)
-                        with o_.skip_synchronize():
-                            scaler.step(o_)
-                    else:
-                        scaler.step(o_)
-                scaler.update()
-            else:
-                total_loss.backward()
-                for o_ in optimizer.values():
-                    o_.step()
-        else:
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                if args.horovod:
-                    optimizer.synchronize()
-                    scaler.unscale_(optimizer)
-                    with optimizer.skip_synchronize():
-                        scaler.step(optimizer)
-                else:
-                    scaler.step(optimizer)
-                scaler.update()
-            else:
-                total_loss.backward()
-                optimizer.step()
-
-        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
-        with torch.no_grad():
-            unwrap_model(model).clap_model.logit_scale_a.clamp_(0, math.log(100))
-            unwrap_model(model).clap_model.logit_scale_t.clamp_(0, math.log(100))
-
-        batch_time_m.update(time.time() - end)
-        end = time.time()
-        batch_count = i + 1
-
-        if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            if isinstance(audio, dict):
-                batch_size = len(audio["waveform"])
-            else:
-                batch_size = len(audio)
-            num_samples = batch_count * batch_size * args.world_size
-            samples_per_epoch = dataloader.num_samples
-            percent_complete = 100.0 * batch_count / num_batches_per_epoch
-
-            # NOTE loss is coarsely sampled, just master node and per log update
-            loss_m.update(total_loss.item(), batch_size)
-            if isinstance(optimizer, dict):
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]}"
-                )
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                }
-            else:
-                logging.info(
-                    f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                    f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                    f"Data (t): {data_time_m.avg:.3f} "
-                    f"Batch (t): {batch_time_m.avg:.3f} "
-                    f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                )
-
-                # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                log_data = {
-                    "loss": loss_m.val,
-                    "data_time": data_time_m.val,
-                    "batch_time": batch_time_m.val,
-                    "lr": optimizer.param_groups[0]["lr"],
-                }
-            for name, val in log_data.items():
-                name = f"train{extra_suffix}/{name}"
-                if tb_writer is not None:
-                    tb_writer.add_scalar(name, val, step)
-                if args.wandb:
-                    assert wandb is not None, "Please install wandb."
-                    wandb.log({name: val, "step": step})
-
-            # resetting batch / data time meters per log window
-            batch_time_m.reset()
-            data_time_m.reset()
-    # end for
-
-
-def evaluate(model, data, epoch, args, tb_writer=None, extra_suffix=""):
-    metrics = {}
-    if not args.parallel_eval:
-        if not is_master(args):
-            return metrics
-    device = torch.device(args.device)
-    model.eval()
-
-    # CHANGE
-    # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
-    # metrics.update(zero_shot_metrics)
-    if is_master(args):
-        print("Evaluating...")
-        metric_names = args.lp_metrics.split(",")
-        eval_tool = LPMetrics(metric_names=metric_names)
-
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    if "val" in data and (
-        args.val_frequency
-        and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
-    ):
-        if args.parallel_eval:
-            dataloader, sampler = data["val"].dataloader, data["val"].sampler
-            if args.distributed and sampler is not None:
-                sampler.set_epoch(epoch)
-            samples_per_val = dataloader.num_samples
-        else:
-            dataloader = data["val"].dataloader
-            num_samples = 0
-            samples_per_val = dataloader.num_samples
-
-        eval_info = {"pred": [], "target": []}
-        with torch.no_grad():
-            for i, batch in enumerate(dataloader):
-                audio = batch  # contains mel_spec, wavform, and longer list
-                class_label = batch["class_label"]
-
-                # audio = audio.to(device=device, non_blocking=True)
-                class_label = class_label.to(device=device, non_blocking=True)
-
-                with autocast():
-                    pred = model(audio, device=device)
-                    if args.parallel_eval:
-                        pred, class_label = lp_gather_features(
-                            pred, class_label, args.world_size, args.horovod
-                        )
-                    eval_info["pred"].append(pred)
-                    eval_info["target"].append(class_label)
-
-                num_samples += class_label.shape[0]
-
-                if (i % 100) == 0:  # and i != 0:
-                    logging.info(
-                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
-                    )
-
-            if is_master(args):
-                eval_info["pred"] = torch.cat(eval_info["pred"], 0).cpu()
-                eval_info["target"] = torch.cat(eval_info["target"], 0).cpu()
-                metric_dict = eval_tool.evaluate_mertics(
-                    eval_info["pred"], eval_info["target"]
-                )
-                metrics.update(metric_dict)
-                if "epoch" not in metrics.keys():
-                    metrics.update({"epoch": epoch})
-
-    if is_master(args):
-        if not metrics:
-            return metrics
-
-        logging.info(
-            f"Eval Epoch: {epoch} "
-            + "\n".join(
-                ["\t".join([f"{m}: {round(metrics[m], 4):.4f}"]) for m in metrics]
-            )
-        )
-        if args.save_logs:
-            for name, val in metrics.items():
-                if tb_writer is not None:
-                    tb_writer.add_scalar(f"val{extra_suffix}/{name}", val, epoch)
-
-            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
-                f.write(json.dumps(metrics))
-                f.write("\n")
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val{extra_suffix}/{name}": val, "epoch": epoch})
-
-        return metrics
-    else:
-        return metrics

models/CLAP/training/main.py

@@ -1,596 +0,0 @@
-from inspect import getargs
-import logging
-import os
-import random
-from datetime import datetime
-import bisect
-import copy
-import numpy as np
-import torch
-import torch.backends.cudnn as cudnn
-from torch import optim
-from torch.cuda.amp import GradScaler
-import faulthandler
-import pathlib
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-try:
-    import torch.utils.tensorboard as tensorboard
-except ImportError:
-    tensorboard = None
-
-try:
-    import horovod.torch as hvd
-except ImportError:
-    hvd = None
-
-from open_clip import create_model_and_transforms, trace_model, create_model
-from training.data import get_data
-from training.distributed import is_master, init_distributed_device, world_info_from_env
-from training.logger import setup_logging
-from training.params import parse_args
-from training.scheduler import cosine_lr
-from training.train import train_one_epoch, evaluate
-from open_clip.utils import dataset_split, get_optimizer
-
-
-def maintain_ckpts(args, startidx, all_idx_len):
-    for i in reversed(range(startidx, all_idx_len)):
-        if os.path.exists(os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt")):
-            os.rename(
-                os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                os.path.join(args.checkpoint_path, f"epoch_top_{i+1}.pt"),
-            )
-    if os.path.exists(
-        os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt")
-    ):
-        os.remove(os.path.join(args.checkpoint_path, f"epoch_top_{all_idx_len}.pt"))
-    return
-
-
-def update_top_k_performance(
-    new_metrics_inputs, current_top_k_ckpt_metrics, args, ckpt, bignumbetter=True
-):
-    """
-    Record the top-k performance of the current epoch.
-    current_top_k_metrics is a dictionary of the form: {1: top_1_ckpt_measure, 2: top_2_ckpt_measure, ...}
-    """
-    if isinstance(new_metrics_inputs, (list, tuple)):
-        new_metrics_inputs = np.mean(new_metrics_inputs)
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, dict):
-        new_metrics_inputs = np.mean(list(new_metrics_inputs.values()))
-        return update_top_k_performance(
-            new_metrics_inputs,
-            current_top_k_ckpt_metrics,
-            args=args,
-            ckpt=ckpt,
-            bignumbetter=bignumbetter,
-        )
-    elif isinstance(new_metrics_inputs, (float, int)):
-        update_flag = {k: False for k in current_top_k_ckpt_metrics.keys()}
-        sorted_keys = sorted(current_top_k_ckpt_metrics.keys())
-        sorted_values = sorted(
-            current_top_k_ckpt_metrics.values(), reverse=bignumbetter
-        )
-        sorted_values_ = copy.deepcopy(sorted_values)
-        sorted_values.append(new_metrics_inputs)
-        sorted_values = sorted(sorted_values, reverse=bignumbetter)
-        sorted_values = sorted_values[:-1]
-
-        if sorted_values == sorted_values_:
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-        else:
-            for i in range(len(sorted_keys)):
-                if current_top_k_ckpt_metrics[sorted_keys[i]] != sorted_values[i]:
-                    current_top_k_ckpt_metrics[sorted_keys[i]] = sorted_values[i]
-                    update_flag[sorted_keys[i]] = True
-            for i in range(len(update_flag)):
-                if update_flag[i]:
-                    maintain_ckpts(args, i, len(sorted_keys))
-                    torch.save(
-                        ckpt,
-                        os.path.join(args.checkpoint_path, f"epoch_top_{i}.pt"),
-                    )
-                    break
-            return current_top_k_ckpt_metrics, new_metrics_inputs
-
-
-# def updateifNone(a, b):
-#     a = b if None else a
-#     return a
-
-
-def is_pretrained_params(n):
-    return (
-        n.startswith("transformer")
-        or n in ["positional_embedding", "text_projection"]
-        or n.startswith("token_embedding")
-        or n.startswith("ln_final")
-        or n.startswith("logit_scale_t")
-    )
-
-
-def random_seed(seed=42, rank=0):
-    torch.manual_seed(seed + rank)
-    np.random.seed(seed + rank)
-    random.seed(seed + rank)
-
-
-def main():
-    args = parse_args()
-    # sanitize model name for filesystem / uri use, easier if we don't use / in name as a rule?
-    args.amodel = args.amodel.replace("/", "-")
-    # download sizes.json file
-
-    # (yusong): the below two lines are for debug
-    # print("setting up faulthandler")
-    # faulthandler.register(10)
-
-    random.seed(args.seed)
-    torch.manual_seed(args.seed)
-    torch.cuda.manual_seed(args.seed)
-    torch.cuda.manual_seed_all(args.seed)
-    np.random.seed(args.seed)
-    if args.tmodel == "bert" or args.tmodel == "roberta" or args.tmodel == "bart":
-        assert (
-            args.pretrained == "" or args.pretrained is None
-        ), "bert/roberta/bart text encoder does not support pretrained models."
-
-    # get the name of the experiments
-    if args.name is None:
-        args.name = "-".join(
-            [
-                datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
-                f"model_{args.amodel}",
-                f"lr_{args.lr}",
-                f"b_{args.batch_size}",
-                f"j_{args.workers}",
-                f"p_{args.precision}",
-            ]
-        )
-
-    # discover initial world args early so we can log properly
-    args.distributed = False
-    args.local_rank, args.rank, args.world_size = world_info_from_env()
-
-    if args.remotedata and is_master(args):
-        for dataset_name in args.datasetnames:
-            for split in dataset_split[dataset_name]:
-                if not os.path.exists(f"./json_files/{dataset_name}/{split}"):
-                    os.makedirs(f"./json_files/{dataset_name}/{split}")
-                os.system(
-                    f"aws s3 cp s3://s-laion-audio/webdataset_tar/{dataset_name}/{split}/sizes.json ./json_files/{dataset_name}/{split}/sizes.json"
-                )
-
-    args.log_path = None
-    if is_master(args, local=args.log_local):
-        log_base_path = os.path.join(args.logs, args.name)
-        os.makedirs(log_base_path, exist_ok=True)
-        log_filename = f"out-{args.rank}" if args.log_local else "out.log"
-        args.log_path = os.path.join(log_base_path, log_filename)
-        if os.path.exists(args.log_path):
-            print(
-                "Error. Experiment already exists. Use --name {} to specify a new experiment."
-            )
-            return -1
-
-    # Set logger
-    args.log_level = logging.DEBUG if args.debug else logging.INFO
-    setup_logging(args.log_path, args.log_level)
-
-    # fully initialize distributed device environment
-    device = init_distributed_device(args)
-
-    args.wandb = "wandb" in args.report_to or "all" in args.report_to
-    args.tensorboard = "tensorboard" in args.report_to or "all" in args.report_to
-    if is_master(args):
-        args.tensorboard_path = (
-            os.path.join(args.logs, args.name, "tensorboard")
-            if args.tensorboard
-            else ""
-        )
-        args.checkpoint_path = os.path.join(args.logs, args.name, "checkpoints")
-        for dirname in [args.tensorboard_path, args.checkpoint_path]:
-            if dirname:
-                os.makedirs(dirname, exist_ok=True)
-    else:
-        args.tensorboard_path = ""
-        args.checkpoint_path = ""
-
-    if args.copy_codebase:
-        copy_codebase(args)
-
-    assert args.precision in ["amp", "fp16", "fp32"]
-    if args.precision == "fp16":
-        logging.warning(
-            "It is recommended to use AMP mixed-precision instead of FP16. "
-            "FP16 support needs further verification and tuning, especially for train."
-        )
-
-    if args.horovod:
-        logging.info(
-            f"Running in horovod mode with multiple processes / nodes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    elif args.distributed:
-        logging.info(
-            f"Running in distributed mode with multiple processes. Device: {args.device}."
-            f"Process (global: {args.rank}, local {args.local_rank}), total {args.world_size}."
-        )
-    else:
-        logging.info(f"Running with a single process. Device {args.device}.")
-
-    logging.info(f"openai cache dir: {os.path.expanduser(args.openai_model_cache_dir)}")
-
-    model, model_cfg = create_model(
-        args.amodel,
-        args.tmodel,
-        args.pretrained,
-        precision=args.precision,
-        device=device,
-        jit=args.torchscript,
-        force_quick_gelu=args.force_quick_gelu,
-        openai_model_cache_dir=os.path.expanduser(args.openai_model_cache_dir),
-        skip_params=True,
-        pretrained_audio=args.pretrained_audio,
-        pretrained_text=args.pretrained_text,
-        enable_fusion=args.enable_fusion,
-        fusion_type=args.fusion_type,
-    )
-
-    if args.horovod:
-        with torch.no_grad():
-            for param in model.parameters():
-                param.set_(param.contiguous())
-
-    if args.trace:
-        model = trace_model(model, batch_size=args.batch_size, device=device)
-
-    if is_master(args):
-        logging.info("Model:")
-        logging.info(f"{str(model)}")
-        logging.info("Params:")
-        params_file = os.path.join(args.logs, args.name, "params.txt")
-        with open(params_file, "w") as f:
-            for name in sorted(vars(args)):
-                val = getattr(args, name)
-                logging.info(f"  {name}: {val}")
-                f.write(f"{name}: {val}\n")
-
-    if args.distributed and not args.horovod:
-        if args.use_bn_sync:
-            model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
-        ddp_args = {}
-        if args.ddp_static_graph:
-            # this doesn't exist in older PyTorch, arg only added if enabled
-            ddp_args["static_graph"] = True
-        model = torch.nn.parallel.DistributedDataParallel(
-            model, device_ids=[device], find_unused_parameters=True, **ddp_args
-        )
-
-    data = get_data(args, model_cfg)
-    assert len(data), "At least one train or eval dataset must be specified."
-    if args.trace:
-        assert "train" not in data, "Cannot train with traced model"
-
-    exclude = (
-        lambda n, p: p.ndim < 2
-        or "bn" in n
-        or "ln" in n
-        or "bias" in n
-        or "logit_scale" in n
-    )
-    include = lambda n, p: not exclude(n, p)
-
-    named_parameters = list(model.named_parameters())
-
-    # freeze text encoder
-    text_freeze_parameters = [p for n, p in named_parameters if "text_branch" in n]
-
-    if args.freeze_text:
-        print("Freeze Text!!!!")
-        for k in text_freeze_parameters:
-            k.requires_grad = False
-
-    gain_or_bias_params = [
-        p for n, p in named_parameters if exclude(n, p) and p.requires_grad
-    ]
-    rest_params = [p for n, p in named_parameters if include(n, p) and p.requires_grad]
-
-    # set wd-related params to 0 if use adam optimizer
-    if args.optimizer == "adam":
-        args.wd = 0
-        args.wd_pretrained = 0
-        args.wd_new = 0
-
-    if args.train_data is None:
-        optimizer = None
-        scheduler = None
-    else:
-        total_steps = data["train"].dataloader.num_batches * args.epochs
-
-        if args.split_opt:
-            for x in ["lr", "beta1", "beta2", "eps", "wd"]:
-                for y in ["_new", "_pretrained"]:
-                    if getattr(args, x + y) is None:
-                        setattr(args, x + y, getattr(args, x))
-
-            gain_or_bias_pretrained_params = [
-                p
-                for n, p in named_parameters
-                if (exclude(n, p) and p.requires_grad) and is_pretrained_params(n)
-            ]
-            rest_pretrained_params = [
-                p
-                for n, p in named_parameters
-                if (include(n, p) and p.requires_grad) and is_pretrained_params(n)
-            ]
-            gain_or_bias_new_params = [
-                p
-                for n, p in named_parameters
-                if (exclude(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-            ]
-            rest_new_params = [
-                p
-                for n, p in named_parameters
-                if (include(n, p) and p.requires_grad) and (not is_pretrained_params(n))
-            ]
-            pretrained_params_optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_pretrained_params, "weight_decay": 0.0},
-                    {
-                        "params": rest_pretrained_params,
-                        "weight_decay": args.wd_pretrained,
-                    },
-                ],
-                lr=args.lr_pretrained,
-                betas=(args.beta1_pretrained, args.beta2_pretrained),
-                eps=args.eps_pretrained,
-                momentum=args.momentum_pretrained,
-                optimizer_name=args.optimizer,
-            )
-            pretrained_params_scheduler = cosine_lr(
-                pretrained_params_optimizer,
-                args.lr_pretrained,
-                args.warmup,
-                total_steps,
-            )
-            new_params_optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_new_params, "weight_decay": 0.0},
-                    {"params": rest_new_params, "weight_decay": args.wd_new},
-                ],
-                lr=args.lr_new,
-                betas=(args.beta1_new, args.beta2_new),
-                eps=args.eps_new,
-                momentum=args.momentum_new,
-                optimizer_name=args.optimizer,
-            )
-
-            new_params_scheduler = cosine_lr(
-                new_params_optimizer, args.lr_new, args.warmup, total_steps
-            )
-
-            optimizer = {
-                "pretrained": pretrained_params_optimizer,
-                "new": new_params_optimizer,
-            }
-            scheduler = {
-                "pretrained": pretrained_params_scheduler,
-                "new": new_params_scheduler,
-            }
-
-            if args.horovod:
-                pretrained_params_optimizer = hvd.DistributedOptimizer(
-                    pretrained_params_optimizer,
-                    named_parameters=model.named_parameters(),
-                )
-                new_params_optimizer = hvd.DistributedOptimizer(
-                    new_params_optimizer, named_parameters=model.named_parameters()
-                )
-                hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                hvd.broadcast_optimizer_state(pretrained_params_optimizer, root_rank=0)
-                hvd.broadcast_optimizer_state(new_params_optimizer, root_rank=0)
-        else:
-            optimizer = get_optimizer(
-                [
-                    {"params": gain_or_bias_params, "weight_decay": 0.0},
-                    {"params": rest_params, "weight_decay": args.wd},
-                ],
-                lr=args.lr,
-                betas=(args.beta1, args.beta2),
-                eps=args.eps,
-                momentum=args.momentum,
-                optimizer_name=args.optimizer,
-            )
-
-            scheduler = cosine_lr(optimizer, args.lr, args.warmup, total_steps)
-
-            if args.horovod:
-                optimizer = hvd.DistributedOptimizer(
-                    optimizer, named_parameters=model.named_parameters()
-                )
-                hvd.broadcast_parameters(model.state_dict(), root_rank=0)
-                hvd.broadcast_optimizer_state(optimizer, root_rank=0)
-
-    scaler = GradScaler() if args.precision == "amp" else None
-
-    # optionally resume from a checkpoint
-    start_epoch = 0
-    if args.resume is not None:
-        if os.path.isfile(args.resume):
-            checkpoint = torch.load(args.resume, map_location=device)
-            if "epoch" in checkpoint:
-                # resuming a train checkpoint w/ epoch and optimizer state
-                start_epoch = checkpoint["epoch"]
-                sd = checkpoint["state_dict"]
-                if not args.distributed and next(iter(sd.items()))[0].startswith(
-                    "module"
-                ):
-                    sd = {k[len("module.") :]: v for k, v in sd.items()}
-                model.load_state_dict(sd)
-                if args.split_opt:
-                    if optimizer is not None:
-                        for k, o_ in optimizer.items():
-                            o_.load_state_dict(checkpoint[k + "_" + "optimizer"])
-                if optimizer is not None:
-                    optimizer.load_state_dict(checkpoint["optimizer"])
-                if scaler is not None and "scaler" in checkpoint:
-                    scaler.load_state_dict(checkpoint["scaler"])
-                logging.info(
-                    f"=> resuming checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            else:
-                # loading a bare (model only) checkpoint for fine-tune or evaluation
-                model.load_state_dict(checkpoint)
-                logging.info(
-                    f"=> loaded checkpoint '{args.resume}' (epoch {start_epoch})"
-                )
-            if args.freeze_text:
-                print("Freeze Text!!!!")
-                for k in text_freeze_parameters:
-                    k.requires_grad = False
-        else:
-            logging.info("=> no checkpoint found at '{}'".format(args.resume))
-
-    cudnn.benchmark = True
-    cudnn.deterministic = False
-
-    # determine if this worker should save logs and checkpoints. only do so if it is rank == 0
-    args.save_logs = args.logs and args.logs.lower() != "none" and is_master(args)
-    writer = None
-    if args.save_logs and args.tensorboard:
-        assert tensorboard is not None, "Please install tensorboard."
-        writer = tensorboard.SummaryWriter(args.tensorboard_path)
-
-    if args.wandb and is_master(args):
-        assert wandb is not None, "Please install wandb."
-        logging.debug("Starting wandb.")
-        args.train_sz = data["train"].dataloader.num_samples
-        if args.val_data is not None:
-            args.val_sz = data["val"].dataloader.num_samples
-        # you will have to configure this for your project!
-        wandb.init(
-            project="clap",
-            notes=args.wandb_notes,
-            name=args.wandb_notes,
-            tags=[],
-            config=vars(args),
-        )
-        if args.debug:
-            wandb.watch(model, log="all")
-        wandb.save(params_file)
-        logging.debug("Finished loading wandb.")
-
-    if "train" not in data:
-        evaluate(model, data, start_epoch, args, writer)
-        return
-    elif start_epoch == 0 and "val" in data and not args.no_eval:
-        evaluate(model, data, 0, args, writer)
-        #  print(f'rank {args.rank}, Start First Evaluation')#  (yusong): for debug
-    if args.save_top_performance:
-        current_top_k_ckpt_metrics = {
-            i: 0 for i in range(args.save_top_performance)
-        }  # initialize the top-k metric for ckpts to 0
-
-    #  print(f'rank {args.rank}, Start Training') #  (yusong): for debug
-    for epoch in range(start_epoch, args.epochs):
-        # freeze the text param after (include) args.freeze_text_after, this is -1 by default
-        if epoch == args.freeze_text_after:
-            print("Text pretrained parameters are freezed since this epoch.")
-            for k in text_freeze_parameters:
-                k.requires_grad = False
-        if is_master(args):
-            logging.info(f"Start epoch {epoch}")
-
-        train_one_epoch(model, data, epoch, optimizer, scaler, scheduler, args, writer)
-        completed_epoch = epoch + 1
-
-        if (
-            any(v in data for v in ("val", "imagenet-val", "imagenet-v2"))
-            and not args.no_eval
-        ):
-            metrics = evaluate(model, data, completed_epoch, args, writer)
-            if args.save_top_performance:
-                top_k_dataset = args.top_k_checkpoint_select_dataset
-                top_k_metric = args.top_k_checkpoint_select_metric
-                filtered_metrics = [
-                    v
-                    for k, v in metrics.items()
-                    if top_k_metric in k and top_k_dataset in k
-                ]  # check all R@10 metrics (all dataset) and use it to update the ckpt
-        # Saving checkpoints.
-        if args.save_logs:
-            if args.split_opt:
-                opt_dict = {
-                    k + "_" + "optimizer": v.state_dict() for k, v in optimizer.items()
-                }
-            else:
-                opt_dict = {"optimizer": optimizer.state_dict()}
-            checkpoint_dict = {
-                "epoch": completed_epoch,
-                "name": args.name,
-                "state_dict": model.state_dict(),
-            }
-            checkpoint_dict.update(opt_dict)
-            if scaler is not None:
-                checkpoint_dict["scaler"] = scaler.state_dict()
-
-            if completed_epoch == args.epochs or (
-                args.save_frequency > 0 and (completed_epoch % args.save_frequency) == 0
-            ):
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_{completed_epoch}.pt"),
-                )
-            if args.save_most_recent:
-                torch.save(
-                    checkpoint_dict,
-                    os.path.join(args.checkpoint_path, f"epoch_latest.pt"),
-                )
-            if args.save_top_performance and not args.no_eval:
-                update_top_k_performance(
-                    filtered_metrics,
-                    current_top_k_ckpt_metrics,
-                    args,
-                    checkpoint_dict,
-                    bignumbetter=True,
-                )
-
-    if args.wandb and is_master(args):
-        wandb.finish()
-
-
-def copy_codebase(args):
-    from shutil import copytree, ignore_patterns
-
-    new_code_path = os.path.join(args.logs, args.name, "code")
-    if os.path.exists(new_code_path):
-        print(
-            f"Error. Experiment already exists at {new_code_path}. Use --name to specify a new experiment."
-        )
-        return -1
-    print(f"Copying codebase to {new_code_path}")
-    current_code_path = os.path.realpath(__file__)
-    for _ in range(3):
-        current_code_path = os.path.dirname(current_code_path)
-    copytree(
-        current_code_path, new_code_path, ignore=ignore_patterns("log", "logs", "wandb")
-    )
-    print("Done copying code.")
-    return 1
-
-
-if __name__ == "__main__":
-    main()

models/CLAP/training/params.py

@@ -1,563 +0,0 @@
-import argparse
-
-
-def get_default_params(model_name):
-    # Params from paper (https://arxiv.org/pdf/2103.00020.pdf)
-    model_name = model_name.lower()
-    if "vit" in model_name:
-        return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.98, "eps": 1.0e-6}
-    else:
-        return {"lr": 5.0e-4, "beta1": 0.9, "beta2": 0.999, "eps": 1.0e-8}
-
-
-def parse_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--train-data",
-        type=str,
-        default=None,
-        help="Path to h5 filewith training data",
-    )
-    parser.add_argument(
-        "--val-data",
-        type=str,
-        default=None,
-        help="Path to h5 file with validation data",
-    )
-    parser.add_argument(
-        "--freeze-text",
-        default=False,
-        action="store_true",
-        help="if you need to freeze the text encoder, make this True",
-    )
-    parser.add_argument(
-        "--freeze-text-after",
-        type=int,
-        default=-1,
-        help="if you need to freeze the text encoder after (include) epoch x, set this param to x. Set -1 to disable it",
-    )
-    parser.add_argument(
-        "--train-ipc",
-        type=str,
-        default=None,
-        help="Path to npy file of the number of instance per class in training data",
-    )
-    parser.add_argument(
-        "--val-ipc",
-        type=str,
-        default=None,
-        help="Path to npy file of the number of instance per class in validation data",
-    )
-    parser.add_argument(
-        "--train-num-samples",
-        type=int,
-        default=None,
-        help="Number of samples in dataset. Required for webdataset if not available in info file.",
-    )
-    parser.add_argument(
-        "--val-num-samples",
-        type=int,
-        default=None,
-        help="Number of samples in dataset. Useful for webdataset if not available in info file.",
-    )
-    parser.add_argument(
-        "--dataset-type",
-        choices=["webdataset", "csv", "auto", "toy"],
-        default="auto",
-        help="Which type of dataset to process.",
-    )
-    parser.add_argument(
-        "--csv-separator",
-        type=str,
-        default="\t",
-        help="For csv-like datasets, which separator to use.",
-    )
-    parser.add_argument(
-        "--csv-img-key",
-        type=str,
-        default="filepath",
-        help="For csv-like datasets, the name of the key for the image paths.",
-    )
-    parser.add_argument(
-        "--csv-caption-key",
-        type=str,
-        default="title",
-        help="For csv-like datasets, the name of the key for the captions.",
-    )
-    parser.add_argument(
-        "--imagenet-val",
-        type=str,
-        default=None,
-        help="Path to imagenet val set for conducting zero shot evaluation.",
-    )
-    parser.add_argument(
-        "--imagenet-v2",
-        type=str,
-        default=None,
-        help="Path to imagenet v2 for conducting zero shot evaluation.",
-    )
-    parser.add_argument(
-        "--datasetnames",
-        nargs="+",
-        default=None,
-        help="If loading webdataset, spedify the dataset names to load. Can be some of these: Clotho, audioset, audiocaps, BBCSoundEffects",
-    )
-    parser.add_argument(
-        "--full-train-dataset",
-        nargs="+",
-        default=None,
-        help="Which dataset will be trained with all the subsets. (train+test)",
-    )
-    parser.add_argument(
-        "--exclude-eval-dataset",
-        nargs="+",
-        default=None,
-        help="Which dataset will be excluded with evaluation",
-    )
-    parser.add_argument(
-        "--datasetinfos",
-        nargs="+",
-        default=None,
-        help="If loading webdataset, spedify the dataset types to load. Can be some of these: train, test, valid, unbalanced_train, balanced_train, eval",
-    )
-    parser.add_argument(
-        "--dataset-proportion",
-        type=float,
-        default=1.0,
-        help="How much proportion of dataset we want to train.",
-    )
-    parser.add_argument(
-        "--remotedata",
-        default=False,
-        action="store_true",
-        help="if the dataset is remote, set this flag",
-    )
-    parser.add_argument(
-        "--class-label-path",
-        type=str,
-        default=None,
-        help="The path of the class label pickle or csv.",
-    )
-    parser.add_argument(
-        "--datasetpath",
-        type=str,
-        default="/mnt/audio_clip/webdataset_tar",
-        help="The path to the dataset",
-    )
-    parser.add_argument(
-        "--logs",
-        type=str,
-        default="./logs/",
-        help="Where to store tensorboard logs. Use None to avoid storing logs.",
-    )
-    parser.add_argument(
-        "--log-local",
-        action="store_true",
-        default=False,
-        help="log files on local master, otherwise global master only.",
-    )
-    parser.add_argument(
-        "--name",
-        type=str,
-        default=None,
-        help="Optional identifier for the experiment when storing logs. Otherwise use current time.",
-    )
-    parser.add_argument(
-        "--workers", type=int, default=1, help="Number of workers per GPU."
-    )
-    parser.add_argument(
-        "--batch-size", type=int, default=64, help="Batch size per GPU."
-    )
-    parser.add_argument(
-        "--epochs", type=int, default=32, help="Number of epochs to train for."
-    )
-    parser.add_argument("--lr", type=float, default=None, help="Learning rate.")
-    parser.add_argument("--beta1", type=float, default=None, help="Adam beta 1.")
-    parser.add_argument("--beta2", type=float, default=None, help="Adam beta 2.")
-    parser.add_argument("--eps", type=float, default=None, help="Adam epsilon.")
-    parser.add_argument("--momentum", type=float, default=None, help="SGD epsilon.")
-    parser.add_argument("--wd", type=float, default=0.2, help="Weight decay.")
-
-    parser.add_argument(
-        "--split-opt",
-        action="store_true",
-        default=False,
-        help="Use this flag to skip the learning rate decay.",
-    )
-    parser.add_argument(
-        "--lr-pretrained", type=float, default=None, help="Learning rate for text."
-    )
-    parser.add_argument(
-        "--beta1-pretrained", type=float, default=None, help="Adam beta 1 for text."
-    )
-    parser.add_argument(
-        "--beta2-pretrained", type=float, default=None, help="Adam beta 2 for text."
-    )
-    parser.add_argument(
-        "--eps-pretrained", type=float, default=None, help="Adam epsilon for text."
-    )
-    parser.add_argument(
-        "--wd-pretrained", type=float, default=0.2, help="Weight decay for text."
-    )
-    parser.add_argument(
-        "--momentum-pretrained", type=float, default=0.9, help="Momentum for text."
-    )
-    parser.add_argument(
-        "--lr-new", type=float, default=None, help="Learning rate for audio."
-    )
-    parser.add_argument(
-        "--beta1-new", type=float, default=None, help="Adam beta 1 for audio."
-    )
-    parser.add_argument(
-        "--beta2-new", type=float, default=None, help="Adam beta 2 for audio."
-    )
-    parser.add_argument(
-        "--eps-new", type=float, default=None, help="Adam epsilon for audio."
-    )
-    parser.add_argument(
-        "--wd-new", type=float, default=0.2, help="Weight decay for audio."
-    )
-    parser.add_argument(
-        "--momentum-new", type=float, default=0.9, help="Momentum for audio."
-    )
-    parser.add_argument(
-        "--warmup", type=int, default=10000, help="Number of steps to warmup for."
-    )
-    parser.add_argument(
-        "--use-bn-sync",
-        default=False,
-        action="store_true",
-        help="Whether to use batch norm sync.",
-    )
-    parser.add_argument(
-        "--skip-scheduler",
-        action="store_true",
-        default=False,
-        help="Use this flag to skip the learning rate decay.",
-    )
-    parser.add_argument(
-        "--save-frequency", type=int, default=1, help="How often to save checkpoints."
-    )
-    parser.add_argument(
-        "--save-top-performance",
-        type=int,
-        default=0,
-        help="Save the top x performance weights if the value >0",
-    )
-    parser.add_argument(
-        "--save-most-recent",
-        action="store_true",
-        default=False,
-        help="Always save the most recent model trained to epoch_latest.pt.",
-    )
-    parser.add_argument(
-        "--zeroshot-frequency", type=int, default=2, help="How often to run zero shot."
-    )
-    parser.add_argument(
-        "--val-frequency",
-        type=int,
-        default=1,
-        help="How often to run evaluation with val data.",
-    )
-    parser.add_argument(
-        "--resume",
-        default=None,
-        type=str,
-        help="path to latest checkpoint (default: none)",
-    )
-    parser.add_argument(
-        "--precision",
-        choices=["amp", "fp16", "fp32"],
-        default="amp",
-        help="Floating point precision.",
-    )
-    parser.add_argument(
-        "--amodel",
-        type=str,
-        default="RN50",
-        help="Name of the audio backbone to use.",
-    )
-    parser.add_argument(
-        "--tmodel",
-        type=str,
-        default="transformer",
-        help="Name of the text backbone to use. Can be [transformer, bert, roberta, bart]",
-    )
-    parser.add_argument(
-        "--pretrained-audio",
-        default="",
-        type=str,
-        help="Use a pretrained audio model weights for the audio encoder of CLAP",
-    )
-    parser.add_argument(
-        "--pretrained-text",
-        default="",
-        type=str,
-        help="Use a pretrained text model weights for the text encoder of CLAP",
-    )
-    parser.add_argument(
-        "--pretrained",
-        default="",
-        type=str,
-        help="Use a pretrained CLIP model weights with the specified tag or file path.",
-    )
-    parser.add_argument(
-        "--pretrained-image",
-        default=False,
-        action="store_true",
-        help="Load imagenet pretrained weights for image tower backbone if available.",
-    )
-    parser.add_argument(
-        "--lock-image",
-        default=False,
-        action="store_true",
-        help="Lock full image tower by disabling gradients.",
-    )
-    parser.add_argument(
-        "--lock-image-unlocked-groups",
-        type=int,
-        default=0,
-        help="Leave last n image tower layer groups unlocked.",
-    )
-    parser.add_argument(
-        "--lock-image-freeze-bn-stats",
-        default=False,
-        action="store_true",
-        help="Freeze BatchNorm running stats in image tower for any locked layers.",
-    )
-    parser.add_argument(
-        "--local-loss",
-        default=False,
-        action="store_true",
-        help="calculate loss w/ local features @ global (instead of realizing full global @ global matrix)",
-    )
-    parser.add_argument(
-        "--gather-with-grad",
-        default=False,
-        action="store_true",
-        help="enable full distributed gradient for feature gather",
-    )
-    parser.add_argument(
-        "--force-quick-gelu",
-        default=False,
-        action="store_true",
-        help="Force use of QuickGELU activation for non-OpenAI transformer models.",
-    )
-    parser.add_argument(
-        "--torchscript",
-        default=False,
-        action="store_true",
-        help="torch.jit.script the model, also uses jit version of OpenAI models if pretrained=='openai'",
-    )
-    parser.add_argument(
-        "--trace",
-        default=False,
-        action="store_true",
-        help="torch.jit.trace the model for inference / eval only",
-    )
-    # arguments for distributed training
-    parser.add_argument(
-        "--dist-url",
-        default="env://",
-        type=str,
-        help="url used to set up distributed training",
-    )
-    parser.add_argument(
-        "--dist-backend", default="nccl", type=str, help="distributed backend"
-    )
-    parser.add_argument(
-        "--report-to",
-        default="",
-        type=str,
-        help="Options are ['wandb', 'tensorboard', 'wandb,tensorboard']",
-    )
-    parser.add_argument(
-        "--wandb-notes", default="", type=str, help="Notes if logging with wandb"
-    )
-    parser.add_argument(
-        "--C", type=float, default=3.16, help="inverse regularizer for logistic reg."
-    )
-    parser.add_argument(
-        "--debug",
-        default=False,
-        action="store_true",
-        help="If true, more information is logged.",
-    )
-    parser.add_argument(
-        "--copy-codebase",
-        default=False,
-        action="store_true",
-        help="If true, we copy the entire base on the log diretory, and execute from there.",
-    )
-    parser.add_argument(
-        "--horovod",
-        default=False,
-        action="store_true",
-        help="Use horovod for distributed training.",
-    )
-    parser.add_argument(
-        "--ddp-static-graph",
-        default=False,
-        action="store_true",
-        help="Enable static graph optimization for DDP in PyTorch >= 1.11.",
-    )
-    parser.add_argument(
-        "--no-set-device-rank",
-        default=False,
-        action="store_true",
-        help="Don't set device index from local rank (when CUDA_VISIBLE_DEVICES restricted to one per proc).",
-    )
-    parser.add_argument("--seed", type=int, default=4242, help="Default random seed.")
-
-    parser.add_argument(
-        "--top-k-checkpoint-select-dataset",
-        type=str,
-        default="all",
-        help="The dataset of selecting top-k checkpoint.",
-    )
-
-    # @R10, @R@5, @R1, mAP@10
-    parser.add_argument(
-        "--top-k-checkpoint-select-metric",
-        type=str,
-        default="_R@10",
-        help="The metric for selecting top-k checkpoint.",
-    )
-    parser.add_argument(
-        "--openai-model-cache-dir",
-        type=str,
-        default="~/.cache/clip",
-        help="Directory to download OpenAI models.",
-    )
-    parser.add_argument(
-        "--optimizer",
-        type=str,
-        default="adamw",
-        help="can be AdamW or SGD",
-    )
-    parser.add_argument(
-        "--parallel-eval",
-        default=False,
-        action="store_true",
-        help="Eval in parallel (multi-GPU, multi-node).",
-    )
-
-    parser.add_argument(
-        "--no-eval",
-        default=False,
-        action="store_true",
-        help="Training without evaluation.",
-    )
-
-    parser.add_argument(
-        "--lp-mlp",
-        default=False,
-        action="store_true",
-        help="Linear Probe using MLP layer or not.",
-    )
-
-    parser.add_argument(
-        "--lp-freeze",
-        default=False,
-        action="store_true",
-        help="Linear Probe using Freeze CLAP or not",
-    )
-
-    parser.add_argument(
-        "--lp-act",
-        default="None",
-        type=str,
-        help="Options are ['relu','elu','prelu','softmax','sigmoid']",
-    )
-
-    parser.add_argument(
-        "--lp-loss", type=str, default="bce", help="Loss func of Linear Probe."
-    )
-
-    parser.add_argument(
-        "--lp-metrics",
-        type=str,
-        default="map,mauc,acc",
-        help="Metrics of Linear Probe.",
-    )
-
-    parser.add_argument(
-        "--lp-lr", type=float, default=1e-4, help="learning rate of linear probe"
-    )
-    parser.add_argument(
-        "--kappa",
-        type=float,
-        default=0,
-        help="the kappa in the weighted contrastive loss, default is to turn off the weighted contrastive loss",
-    )
-
-    parser.add_argument(
-        "--data-filling",
-        type=str,
-        default="pad",
-        help="type of data filling when the audio length is shorter than the max length."
-        "Can be one of the following: repeat, repeatpad, pad",
-    )
-    parser.add_argument(
-        "--data-truncating",
-        type=str,
-        default="rand_trunc",
-        help="type of data truncation when the audio length is longer than the max length."
-        "Can be one of the following: rand_trunc, fusion",
-    )
-
-    parser.add_argument(
-        "--clap-mlploss",
-        default=False,
-        action="store_true",
-        help="Using MLP loss for CLAP model or not",
-    )
-
-    parser.add_argument(
-        "--wandb-id",
-        type=str,
-        default=None,
-        help="the id of wandb experiment to restore.",
-    )
-
-    parser.add_argument(
-        "--sleep", type=float, default=0, help="sleep n seconds before start training"
-    )
-
-    # variable length processing
-    parser.add_argument(
-        "--enable-fusion",
-        default=False,
-        action="store_true",
-        help="Enable feature funsion for variable-length data",
-    )
-
-    parser.add_argument(
-        "--fusion-type",
-        type=str,
-        default="None",
-        help="Type is among ['channel_map', 'daf_1d','aff_1d','iaff_1d','daf_2d','aff_2d','iaff_2d']",
-    )
-
-    parser.add_argument(
-        "--mixup",
-        default=False,
-        action="store_true",
-        help="Enable mixup in finetuning training.",
-    )
-    parser.add_argument(
-        "--text-augment-selection",
-        type=str,
-        default=None,
-        help="For selecting levels of augmented text. Type is among ['all', 'augment_only', 'none']",
-    )
-
-    args = parser.parse_args()
-
-    # If some params are not passed, we use the default values based on model name.
-    default_params = get_default_params(args.amodel)
-    for name, val in default_params.items():
-        if getattr(args, name) is None:
-            setattr(args, name, val)
-
-    return args

models/CLAP/training/scheduler.py

@@ -1,24 +0,0 @@
-import numpy as np
-
-
-def assign_learning_rate(optimizer, new_lr):
-    for param_group in optimizer.param_groups:
-        param_group["lr"] = new_lr
-
-
-def _warmup_lr(base_lr, warmup_length, step):
-    return base_lr * (step + 1) / warmup_length
-
-
-def cosine_lr(optimizer, base_lr, warmup_length, steps):
-    def _lr_adjuster(step):
-        if step < warmup_length:
-            lr = _warmup_lr(base_lr, warmup_length, step)
-        else:
-            e = step - warmup_length
-            es = steps - warmup_length
-            lr = 0.5 * (1 + np.cos(np.pi * e / es)) * base_lr
-        assign_learning_rate(optimizer, lr)
-        return lr
-
-    return _lr_adjuster

models/CLAP/training/train.py

@@ -1,838 +0,0 @@
-import json
-import logging
-import math
-import os
-import time
-from contextlib import suppress
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-try:
-    import wandb
-except ImportError:
-    wandb = None
-
-from open_clip import ClipLoss, gather_features
-from .distributed import is_master
-from .zero_shot import zero_shot_eval
-
-
-class AverageMeter(object):
-    """Computes and stores the average and current value"""
-
-    def __init__(self):
-        self.reset()
-
-    def reset(self):
-        self.val = 0
-        self.avg = 0
-        self.sum = 0
-        self.count = 0
-
-    def update(self, val, n=1):
-        self.val = val
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
-
-
-def unwrap_model(model):
-    if hasattr(model, "module"):
-        return model.module
-    else:
-        return model
-
-
-def train_one_epoch(
-    model, data, epoch, optimizer, scaler, scheduler, args, tb_writer=None
-):
-    device = torch.device(args.device)
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    model.train()
-    loss = ClipLoss(
-        local_loss=args.local_loss,
-        gather_with_grad=args.gather_with_grad,
-        cache_labels=True,
-        rank=args.rank,
-        world_size=args.world_size,
-        use_horovod=args.horovod,
-        mlp_loss=args.clap_mlploss,
-        weight_loss_kappa=args.kappa,
-    )
-
-    dataloader, sampler = data["train"].dataloader, data["train"].sampler
-    if args.distributed and sampler is not None:
-        sampler.set_epoch(epoch)
-    num_batches_per_epoch = dataloader.num_batches
-    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))
-
-    # for toy dataset
-    if args.dataset_type == "toy":
-        dataloader.dataset.generate_queue()
-
-    loss_m = AverageMeter()
-    batch_time_m = AverageMeter()
-    data_time_m = AverageMeter()
-    end = time.time()
-
-    for i, batch in enumerate(dataloader):
-        # logging.info(f"batch {i} of {num_batches_per_epoch}")
-        step = num_batches_per_epoch * epoch + i
-        if isinstance(scheduler, dict):
-            for s in scheduler.values():
-                s(step)
-        else:
-            scheduler(step)
-        audios = batch  # contains mel_spec, wavform, and longer list
-        texts = batch["text"]
-        # audios = audios.to(device=device, non_blocking=True)
-        # texts = texts.to(device=device, non_blocking=True)
-
-        data_time_m.update(time.time() - end)
-        if isinstance(optimizer, dict):
-            for o_ in optimizer.values():
-                o_.zero_grad()
-        else:
-            optimizer.zero_grad()
-
-        with autocast():
-            (
-                audio_features,
-                text_features,
-                audio_features_mlp,
-                text_features_mlp,
-                logit_scale_a,
-                logit_scale_t,
-            ) = model(audios, texts, device)
-
-            if args.clap_mlploss:
-                total_loss = loss(
-                    audio_features=audio_features,
-                    text_features=text_features,
-                    logit_scale_a=logit_scale_a,
-                    logit_scale_t=logit_scale_t,
-                    audio_features_mlp=audio_features_mlp,
-                    text_features_mlp=text_features_mlp,
-                )
-            else:
-                total_loss = loss(
-                    audio_features=audio_features,
-                    text_features=text_features,
-                    logit_scale_a=logit_scale_a,
-                )
-        if isinstance(optimizer, dict):
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                for o_ in optimizer.values():
-                    if args.horovod:
-                        o_.synchronize()
-                        scaler.unscale_(o_)
-                        with o_.skip_synchronize():
-                            scaler.step(o_)
-                    else:
-                        scaler.step(o_)
-                scaler.update()
-            else:
-                total_loss.backward()
-                for o_ in optimizer.values():
-                    o_.step()
-        else:
-            if scaler is not None:
-                scaler.scale(total_loss).backward()
-                if args.horovod:
-                    optimizer.synchronize()
-                    scaler.unscale_(optimizer)
-                    with optimizer.skip_synchronize():
-                        scaler.step(optimizer)
-                else:
-                    scaler.step(optimizer)
-                scaler.update()
-            else:
-                total_loss.backward()
-                optimizer.step()
-
-        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
-        with torch.no_grad():
-            unwrap_model(model).logit_scale_a.clamp_(0, math.log(100))
-            if args.clap_mlploss:
-                unwrap_model(model).logit_scale_t.clamp_(0, math.log(100))
-
-        batch_time_m.update(time.time() - end)
-        end = time.time()
-        batch_count = i + 1
-        if is_master(args) and (i % 100 == 0 or batch_count == num_batches_per_epoch):
-            if isinstance(audios, dict):
-                batch_size = len(audios["waveform"])
-            else:
-                batch_size = len(audios)
-            num_samples = batch_count * batch_size * args.world_size
-            samples_per_epoch = dataloader.num_samples
-            percent_complete = 100.0 * batch_count / num_batches_per_epoch
-
-            # NOTE loss is coarsely sampled, just master node and per log update
-            loss_m.update(total_loss.item(), batch_size)
-            logit_scale_scalar_a = logit_scale_a.item()
-            logit_scale_scalar_t = logit_scale_t.item()
-            if isinstance(optimizer, dict):
-                if args.clap_mlploss:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                    )
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "scale_text": logit_scale_scalar_t,
-                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                    }
-                else:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {[o_.param_groups[0]['lr'] for o_ in optimizer.values()]} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    )
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "lr": [o_.param_groups[0]["lr"] for o_ in optimizer.values()],
-                    }
-
-            else:
-                if args.clap_mlploss:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                        f"Logit Scale Text: {logit_scale_scalar_t:.3f}"
-                    )
-
-                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "scale_text": logit_scale_scalar_t,
-                        "lr": optimizer.param_groups[0]["lr"],
-                    }
-                else:
-                    logging.info(
-                        f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
-                        f"Loss: {loss_m.val:#.5g} ({loss_m.avg:#.4g}) "
-                        f"Data (t): {data_time_m.avg:.3f} "
-                        f"Batch (t): {batch_time_m.avg:.3f} "
-                        f"LR: {optimizer.param_groups[0]['lr']:5f} "
-                        f"Logit Scale Audio: {logit_scale_scalar_a:.3f}"
-                    )
-
-                    # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
-                    log_data = {
-                        "loss": loss_m.val,
-                        "data_time": data_time_m.val,
-                        "batch_time": batch_time_m.val,
-                        "scale_audio": logit_scale_scalar_a,
-                        "lr": optimizer.param_groups[0]["lr"],
-                    }
-            for name, val in log_data.items():
-                name = "train/" + name
-                if tb_writer is not None:
-                    tb_writer.add_scalar(name, val, step)
-                if args.wandb:
-                    assert wandb is not None, "Please install wandb."
-                    wandb.log({name: val, "step": step})
-
-            # resetting batch / data time meters per log window
-            batch_time_m.reset()
-            data_time_m.reset()
-    # end for
-
-
-def evaluate(model, data, epoch, args, tb_writer=None):
-    metrics = {}
-    if not args.parallel_eval:
-        if not is_master(args):
-            return metrics
-    device = torch.device(args.device)
-    model.eval()
-
-    # CHANGE
-    # zero_shot_metrics = zero_shot_eval(model, data, epoch, args)
-    # metrics.update(zero_shot_metrics)
-    if is_master(args):
-        print("Evaluating...")
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    if args.val_dataset_names == ["Clotho", "audiocaps"]:
-        # if only clotho and audiocaps are used, then we will use a different evaluation function.
-        # This is because in the Clotho and audiocaps valid and test set, there are 5 text for 1 audio.
-        if args.parallel_eval:
-            # (yusong): just a hack here. Don't use parallel eval when evaluating only clotho and audiocaps.
-            raise NotImplementedError(
-                "Parallel evaluation not supported for eval only Clotho and audiocaps."
-            )
-        val_metrics_per_dataset = evaluate_clotho_audiocaps(
-            model, data, epoch, args, autocast, device, tb_writer
-        )
-        for m in val_metrics_per_dataset.values():
-            metrics.update(m)
-        if "epoch" not in metrics.keys():
-            metrics.update({"epoch": epoch})
-        metrics = select_top_metric_clotho_audiocaps(
-            metrics, val_metrics_per_dataset, args
-        )
-    elif "val" in data and (
-        args.val_frequency
-        and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)
-    ):
-        dataloader = data["val"].dataloader
-        num_samples = 0
-        samples_per_val = dataloader.num_samples
-
-        # FIXME this does not scale past small eval datasets
-        # all_audio_features @ all_text_features will blow up memory and compute very quickly
-        eval_info = {}
-        if args.clap_mlploss:
-            eval_info["all"] = {
-                "cumulative_loss": 0.0,
-                "num_samples": 0,
-                "all_audio_features": [],
-                "all_text_features": [],
-                "all_audio_features_mlp": [],
-                "all_text_features_mlp": [],
-            }  # cumulative_loss = 0.0
-        else:
-            eval_info["all"] = {
-                "cumulative_loss": 0.0,
-                "num_samples": 0,
-                "all_audio_features": [],
-                "all_text_features": [],
-            }  # cumu
-        # all_audio_features, all_text_features, all_audio_features_mlp, all_text_features_mlp = [], [], [], []
-        with torch.no_grad():
-            for i, batch in enumerate(dataloader):
-                audios = batch  # contains mel_spec, wavform, and longer list
-                texts = batch["text"]
-                # audios = audios.to(device=device, non_blocking=True)
-
-                all_names = list(
-                    set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
-                )
-                for name in all_names:
-                    if name not in eval_info.keys():
-                        if args.clap_mlploss:
-                            eval_info[name] = {
-                                "cumulative_loss": 0.0,
-                                "num_samples": 0,
-                                "all_audio_features": [],
-                                "all_text_features": [],
-                                "all_audio_features_mlp": [],
-                                "all_text_features_mlp": [],
-                            }
-                        else:
-                            eval_info[name] = {
-                                "cumulative_loss": 0.0,
-                                "num_samples": 0,
-                                "all_audio_features": [],
-                                "all_text_features": [],
-                            }
-                with autocast():
-                    (
-                        audio_features,
-                        text_features,
-                        audio_features_mlp,
-                        text_features_mlp,
-                        logit_scale_a,
-                        logit_scale_t,
-                    ) = model(audios, texts, device)
-
-                    if args.parallel_eval:
-                        # multi-GPU eval
-                        if args.clap_mlploss:
-                            (
-                                audio_features,
-                                text_features,
-                                audio_features_mlp,
-                                text_features_mlp,
-                            ) = gather_features(
-                                audio_features=audio_features,
-                                text_features=text_features,
-                                audio_features_mlp=audio_features_mlp,
-                                text_features_mlp=text_features_mlp,
-                                local_loss=False,
-                                gather_with_grad=False,
-                                rank=args.rank,
-                                world_size=args.world_size,
-                                use_horovod=args.horovod,
-                                mlp_loss=args.clap_mlploss,
-                            )
-                        else:
-                            (audio_features, text_features,) = gather_features(
-                                audio_features=audio_features,
-                                text_features=text_features,
-                                local_loss=False,
-                                gather_with_grad=False,
-                                rank=args.rank,
-                                world_size=args.world_size,
-                                use_horovod=args.horovod,
-                                mlp_loss=args.clap_mlploss,
-                            )
-
-                    if is_master(args):
-                        num_samples += audio_features.shape[0]
-                        for n in [*all_names, "all"]:
-                            if n == "all":
-                                eval_info[n]["all_audio_features"].append(
-                                    audio_features.cpu()
-                                )
-                                eval_info[n]["all_text_features"].append(
-                                    text_features.cpu()
-                                )
-                                if args.clap_mlploss:
-                                    eval_info[n]["all_audio_features_mlp"].append(
-                                        audio_features_mlp.cpu()
-                                    )
-                                    eval_info[n]["all_text_features_mlp"].append(
-                                        text_features_mlp.cpu()
-                                    )
-                            else:
-                                idx = np.where(
-                                    np.array(
-                                        [
-                                            "-".join(b.split("/")[-3:-1])
-                                            for b in batch["__url__"]
-                                        ]
-                                    )
-                                    == n
-                                )[0]
-                                eval_info[n]["all_audio_features"].append(
-                                    audio_features.cpu().index_select(
-                                        0, torch.tensor(idx).long()
-                                    )
-                                )
-                                eval_info[n]["all_text_features"].append(
-                                    text_features.cpu().index_select(
-                                        0, torch.tensor(idx).long()
-                                    )
-                                )
-                                if args.clap_mlploss:
-                                    eval_info[n]["all_audio_features_mlp"].append(
-                                        audio_features_mlp.cpu().index_select(
-                                            0, torch.tensor(idx).long()
-                                        )
-                                    )
-                                    eval_info[n]["all_text_features_mlp"].append(
-                                        text_features_mlp.cpu().index_select(
-                                            0, torch.tensor(idx).long()
-                                        )
-                                    )
-                        #  print(f'eval step {i}') #  (yusong): for debug
-
-                # cumulative_loss += total_loss * batch_size
-                # num_samples += batch_size
-                if is_master(args) and (i % 100) == 0:  # and i != 0:
-                    logging.info(
-                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]"
-                    )
-            if is_master(args):
-                val_metrics_per_dataset = {}
-                for n in eval_info.keys():
-                    if args.clap_mlploss:
-                        metrics_single_dataset = get_metrics(
-                            audio_features=torch.cat(
-                                eval_info[n]["all_audio_features"]
-                            ),
-                            text_features=torch.cat(eval_info[n]["all_text_features"]),
-                            logit_scale_a=logit_scale_a.cpu(),
-                            audio_features_mlp=torch.cat(
-                                eval_info[n]["all_audio_features_mlp"]
-                            ),
-                            text_features_mlp=torch.cat(
-                                eval_info[n]["all_text_features_mlp"]
-                            ),
-                            logit_scale_t=logit_scale_t.cpu(),
-                            mlp_loss=args.clap_mlploss,
-                        )
-                    else:
-                        metrics_single_dataset = get_metrics(
-                            audio_features=torch.cat(
-                                eval_info[n]["all_audio_features"]
-                            ),
-                            text_features=torch.cat(eval_info[n]["all_text_features"]),
-                            logit_scale_a=logit_scale_a.cpu(),
-                            mlp_loss=args.clap_mlploss,
-                        )
-                    val_metrics_per_dataset[n] = {
-                        n + "/" + k: v for k, v in metrics_single_dataset.items()
-                    }
-                    metrics.update(val_metrics_per_dataset[n])
-                    if "epoch" not in metrics.keys():
-                        metrics.update({"epoch": epoch})
-    if is_master(args):
-        if not metrics:
-            return metrics
-
-        logging.info(
-            f"Eval Epoch: {epoch} "
-            + "\n".join(
-                [
-                    "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in m.items()])
-                    for m in val_metrics_per_dataset.values()
-                ]
-            )
-        )
-
-        if args.save_logs:
-            for name, val in metrics.items():
-                if tb_writer is not None:
-                    tb_writer.add_scalar(f"val/{name}", val, epoch)
-
-            with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
-                f.write(json.dumps(metrics))
-                f.write("\n")
-
-        if args.wandb:
-            assert wandb is not None, "Please install wandb."
-            for name, val in metrics.items():
-                wandb.log({f"val/{name}": val, "epoch": epoch})
-
-        return metrics
-    else:
-        return metrics
-
-
-def get_metrics(
-    audio_features,
-    text_features,
-    logit_scale_a,
-    audio_features_mlp=None,
-    text_features_mlp=None,
-    logit_scale_t=None,
-    mlp_loss=False,
-):
-    metrics = {}
-    if mlp_loss:
-        # Set up audio to text & text to audio similary matrice
-        a_logits_per_audio = (
-            (logit_scale_a * audio_features @ text_features_mlp.t()).detach().cpu()
-        )
-        a_logits_per_text = a_logits_per_audio.t().detach().cpu()
-        t_logits_per_audio = (
-            (logit_scale_t * audio_features_mlp @ text_features.t()).detach().cpu()
-        )
-        t_logits_per_text = t_logits_per_audio.t().detach().cpu()
-
-        labels = torch.arange(audio_features.shape[0]).long()
-        # Change the loss from two terms into four terms with 2x2 combined CE loss
-        total_loss = (
-            F.cross_entropy(a_logits_per_audio, labels)
-            + F.cross_entropy(a_logits_per_text, labels)
-            + F.cross_entropy(t_logits_per_audio, labels)
-            + F.cross_entropy(t_logits_per_text, labels)
-        ) / 4
-
-        metrics[f"cumulative_loss"] = total_loss.item()
-        metrics[f"num_samples"] = audio_features.shape[0]
-
-        logits = {
-            "audio_to_text": (a_logits_per_audio + t_logits_per_audio) / 2,
-            "text_to_audio": (a_logits_per_text + t_logits_per_text) / 2,
-        }
-        ground_truth = torch.arange(len(text_features)).view(-1, 1)
-
-    else:
-        # print("text_features", text_features)
-        # print("text_features.shape", text_features.shape)
-        logits_per_audio = (
-            (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
-        )
-        logits_per_text = logits_per_audio.t().detach().cpu()
-
-        labels = torch.arange(audio_features.shape[0]).long()
-        # Change the loss from two terms into four terms with 2x2 combined CE loss
-        total_loss = (
-            F.cross_entropy(logits_per_audio, labels)
-            + F.cross_entropy(logits_per_text, labels)
-        ) / 2
-
-        metrics[f"cumulative_loss"] = total_loss.item()
-        metrics[f"num_samples"] = audio_features.shape[0]
-
-        logits = {"audio_to_text": logits_per_audio, "text_to_audio": logits_per_text}
-
-        ground_truth = torch.arange(len(text_features)).view(-1, 1)
-
-    for name, logit in logits.items():
-        ranking = torch.argsort(logit, descending=True)
-        preds = torch.where(ranking == ground_truth)[
-            1
-        ]  # (yusong) this line is slow because it uses single thread
-        preds = preds.detach().cpu().numpy()
-        metrics[f"{name}_mean_rank"] = preds.mean() + 1
-        metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
-        for k in [1, 5, 10]:
-            metrics[f"{name}_R@{k}"] = np.mean(preds < k)
-        # map@10
-        metrics[f"{name}_mAP@10"] = np.mean(np.where(preds < 10, 1 / (preds + 1), 0.0))
-
-    return metrics
-
-
-def evaluate_clotho_audiocaps(
-    model, data, epoch, args, autocast, device, tb_writer=None
-):
-    """
-    Adapted from https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py.
-    1. for text-to-audio retrieval, do 5 times and average the results
-    2. for R@1, R@5, R@10 in audio-to-text retrieval, take the best rank among 5 text
-    3. for map@10 in audio-to-text retrieval:
-        3.1: sort the rank of 5 text
-        3.2: exclude the rank >=10 (0-index)
-        3.3: compute the map regarding the remaining ranks: np.mean(np.arange(1, len(ranks)+1) / ranks).
-        (3.3) That is, take the top ranks of 5 text that is < 10, and assign the descending number as ground truth.
-        (3.3) E.g.: the ground truth of first rank of the 5 text should be 1, the second rank should be 2, etc.
-    """
-    # TODO: (yusong) only support single GPU evaluation and only support non-mlp case for now.
-    dataloader = data["val"].dataloader
-    with torch.no_grad():
-        eval_info = {}
-        for i, batch in enumerate(dataloader):
-            audios = batch  # contains mel_spec, wavform, and longer list
-
-            # each item in the list has 5 texts
-            if args.tmodel == "transformer":
-                from open_clip import tokenize
-
-                texts = [tokenize(t) for t in batch["full_text"]]
-                texts = torch.cat(texts)
-            else:
-                from .data import tokenizer
-
-                texts = [
-                    tokenizer(t) for t in batch["full_text"]
-                ]  # 5 texts for each audio
-                texts = {
-                    k: torch.cat([t[k] for t in texts]) for k in texts[0].keys()
-                }  # 5 x batch
-
-            # audios = audios.to(device=device, non_blocking=True)
-
-            all_names = list(
-                set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
-            )
-            for name in all_names:
-                if name not in eval_info.keys():
-                    # we will not use mlp outputs even if args.clap_mlploss=True
-                    eval_info[name] = {
-                        "cumulative_loss": 0.0,
-                        "num_samples": 0,
-                        "all_audio_features": [],
-                        "all_text_features": [],
-                    }
-            with autocast():
-                audio_features = model(audios, None, device)
-                text_features = model(None, texts, device)
-                audio_features = F.normalize(audio_features, dim=-1)
-                text_features = F.normalize(text_features, dim=-1)
-
-                all_names = list(
-                    set(["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]])
-                )
-                for n in all_names:
-                    idx = np.where(
-                        np.array(
-                            ["-".join(b.split("/")[-3:-1]) for b in batch["__url__"]]
-                        )
-                        == n
-                    )[0]
-                    eval_info[n]["all_audio_features"].append(
-                        audio_features.cpu().index_select(0, torch.tensor(idx).long())
-                    )
-                    # (yusong) please double-check. This is for selecting 5 text features at once.
-                    # because idx is a list of indices in size of num_samples,
-                    # and text_features is a tensor of size (5*num_samples, dim)
-                    # so we need to select 5 consecutive indices at once for a single index in idx.
-                    eval_info[n]["all_text_features"].append(
-                        text_features.cpu()
-                        .reshape([-1, 5, text_features.shape[1]])
-                        .index_select(0, torch.tensor(idx).long())
-                        .reshape([-1, text_features.shape[1]])
-                    )
-
-        val_metrics_all = {}
-
-        for n in eval_info.keys():
-            logit_scale_a, logit_scale_t = model(None, None, device)
-            logit_scale_a = logit_scale_a.cpu()
-
-            audio_features = torch.cat(eval_info[n]["all_audio_features"], dim=0)
-            text_features = torch.cat(eval_info[n]["all_text_features"], dim=0)
-
-            logits_per_audio = (
-                (logit_scale_a * audio_features @ text_features.t()).detach().cpu()
-            )
-            logits_per_text = logits_per_audio.t().detach().cpu()
-
-            # logits_per_audio shape: [num_samples, num_samples*5]
-            # logits_per_text shape: [num_samples*5, num_samples]
-
-            logging.info(
-                f"dataset {n}, logits_per_audio shape: {logits_per_audio.shape}, "
-                f"logits_per_text shape: {logits_per_text.shape}"
-            )
-
-            metrics = {}
-            num_samples = audio_features.shape[0]
-            metrics[f"num_samples"] = num_samples
-
-            # (yusong) the following code is very important, please double-check:
-            # logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d]
-            # logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
-            # Those two are retrieving one of the 5 text for each audio.
-            labels = torch.arange(audio_features.shape[0]).long()
-            audio_to_text_loss = [
-                F.cross_entropy(
-                    logits_per_audio.reshape(num_samples, num_samples, 5)[:, :, d],
-                    labels,
-                )
-                for d in range(5)
-            ]
-            text_to_audio_loss = [
-                F.cross_entropy(
-                    logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :],
-                    labels,
-                )
-                for d in range(5)
-            ]
-            total_loss = (np.mean(audio_to_text_loss) + np.mean(text_to_audio_loss)) / 2
-
-            metrics[f"cumulative_loss"] = total_loss.item()
-
-            # text to audio: do 5 times
-            pred_text = []
-            for d in range(5):
-                logit = logits_per_text.reshape(num_samples, 5, num_samples)[:, d, :]
-                ground_truth = torch.arange(len(logit)).view(-1, 1)
-                ranking = torch.argsort(
-                    logit, descending=True
-                )  # [num_samples, num_samples]
-                preds = torch.where(ranking == ground_truth)[1]
-                pred_text.append(preds.detach().cpu().numpy())
-            pred_text_concat = np.concatenate(pred_text, axis=0)  # [5*num_samples]
-            metrics[f"text_to_audio_mean_rank"] = pred_text_concat.mean() + 1
-            metrics[f"text_to_audio_median_rank"] = (
-                np.floor(np.median(pred_text_concat)) + 1
-            )
-            for k in [1, 5, 10]:
-                metrics[f"text_to_audio_R@{k}"] = np.mean(pred_text_concat < k)
-            # map@10
-            metrics[f"text_to_audio_mAP@10"] = np.mean(
-                np.where(pred_text_concat < 10, 1 / (pred_text_concat + 1), 0.0)
-            )
-
-            # audio to text: take the best result
-            # for audio to text map 10, sort and assign descending ground truth.
-            # see https://github.com/XinhaoMei/audio-text_retrieval/blob/main/tools/utils.py#L103
-            # map@10
-            map_all = []
-            pred_audio_all = []
-            for d in range(num_samples):
-                # logits_per_audio: [num_samples, num_samples*5]
-                logit_single = logits_per_audio[d, :]  # [5*num_samples]
-                # Ground-truth index: [d*5, d*5+1, d*5+2, d*5+3, d*5+4]
-                ranking = torch.argsort(
-                    logit_single, descending=True
-                )  # [5*num_samples]
-                # ranking: the index of first match, second match, ...
-                ground_truth = torch.arange(d * 5, d * 5 + 5)[None]
-                all_pred = torch.where(
-                    torch.stack([ranking] * 5) == ground_truth.view(-1, 1)
-                )[1]
-                min_pred = torch.min(all_pred)
-                pred_audio_all.append(min_pred.detach().cpu().numpy())
-                all_pred_filter = all_pred[all_pred < 10].detach().cpu().numpy()
-                # /5 because we have 5 text, so it means for the text rank >=10 we count as 0.
-                map_single = (
-                    np.sum(
-                        (np.arange(1, len(all_pred_filter) + 1) / (all_pred_filter + 1))
-                    )
-                    / 5
-                )
-                map_all.append(map_single)
-            metrics[f"audio_to_text_mAP@10"] = np.mean(map_all)
-            for k in [1, 5, 10]:
-                metrics[f"audio_to_text_R@{k}"] = np.mean(np.array(pred_audio_all) < k)
-
-            val_metrics_all[n] = {n + "/" + k: v for k, v in metrics.items()}
-    return val_metrics_all
-
-
-def calculate_selection_performance_clotho_audiocaps(val_metrics_per_dataset):
-    """
-    Calculate performance for Clotho+AudioCaps for model selection.
-    """
-    selection_performance_all = []
-    for n in val_metrics_per_dataset.keys():
-        selection_performance = (
-            val_metrics_per_dataset[n][f"{n}/audio_to_text_mAP@10"]
-            + val_metrics_per_dataset[n][f"{n}/text_to_audio_mAP@10"]
-        ) / 2
-        selection_performance_all.append(selection_performance)
-    return np.mean(selection_performance_all)
-
-
-def select_top_metric_clotho_audiocaps(metrics, val_metrics_per_dataset, args):
-    # val_metrics_per_dataset: dict, key: dataset name, value: dict, key: metric name, value: metric value
-    # metrics: dict, key: metric name, value: metric value
-    # Hack: use args to save the top performance
-    if not hasattr(args, "top_selection_performance"):
-        selection_performance = calculate_selection_performance_clotho_audiocaps(
-            val_metrics_per_dataset
-        )
-        # TODO: write the if and else together
-        metric_update = {}
-        for n in val_metrics_per_dataset.keys():
-            for k in val_metrics_per_dataset[n].keys():
-                metric_update[
-                    k.split("/")[0] + "-top" + "/" + k.split("/")[1]
-                ] = val_metrics_per_dataset[n][k]
-        metric_update["top_selection_performance"] = selection_performance
-        metric_update["top-selection-epoch"] = metrics["epoch"]
-        metrics.update(metric_update)
-        args.top_metric = metric_update
-        args.top_selection_performance = selection_performance
-    else:
-        selection_performance_new = calculate_selection_performance_clotho_audiocaps(
-            val_metrics_per_dataset
-        )
-        selection_performance_old = args.top_selection_performance
-        if selection_performance_new > selection_performance_old:
-            metric_update = {}
-            for n in val_metrics_per_dataset.keys():
-                for k in val_metrics_per_dataset[n].keys():
-                    metric_update[
-                        k.split("/")[0] + "-top" + "/" + k.split("/")[1]
-                    ] = val_metrics_per_dataset[n][k]
-            metric_update["top_selection_performance"] = selection_performance_new
-            metric_update["top-selection-epoch"] = metrics["epoch"]
-            metrics.update(metric_update)
-            args.top_metric = metric_update
-            args.top_selection_performance = selection_performance_new
-        else:
-            metrics.update(args.top_metric)
-    return metrics

models/CLAP/training/zero_shot.py

@@ -1,95 +0,0 @@
-# NOTE: This script is currently not supported for CLAP.
-import logging
-from contextlib import suppress
-
-import torch
-import torch.nn.functional as F
-from tqdm import tqdm
-
-from open_clip import tokenize
-from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
-
-
-def zero_shot_classifier(model, classnames, templates, args):
-    with torch.no_grad():
-        zeroshot_weights = []
-        for classname in tqdm(classnames):
-            texts = [template(classname) for template in templates]  # format with class
-            texts = tokenize(texts).to(args.device)  # tokenize
-            if args.distributed and not args.horovod:
-                class_embeddings = model.module.encode_text(texts)
-            else:
-                class_embeddings = model.encode_text(texts)
-            class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
-            class_embedding /= class_embedding.norm()
-            zeroshot_weights.append(class_embedding)
-        zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.device)
-    return zeroshot_weights
-
-
-def accuracy(output, target, topk=(1,)):
-    pred = output.topk(max(topk), 1, True, True)[1].t()
-    correct = pred.eq(target.view(1, -1).expand_as(pred))
-    return [
-        float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy())
-        for k in topk
-    ]
-
-
-def run(model, classifier, dataloader, args):
-    autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
-    with torch.no_grad():
-        top1, top5, n = 0.0, 0.0, 0.0
-        for images, target in tqdm(dataloader, unit_scale=args.batch_size):
-            images = images.to(args.device)
-            target = target.to(args.device)
-
-            with autocast():
-                # predict
-                if args.distributed and not args.horovod:
-                    image_features = model.module.encode_image(images)
-                else:
-                    image_features = model.encode_image(images)
-                image_features = F.normalize(image_features, dim=-1)
-                logits = 100.0 * image_features @ classifier
-
-            # measure accuracy
-            acc1, acc5 = accuracy(logits, target, topk=(1, 5))
-            top1 += acc1
-            top5 += acc5
-            n += images.size(0)
-
-    top1 = top1 / n
-    top5 = top5 / n
-    return top1, top5
-
-
-def zero_shot_eval(model, data, epoch, args):
-    if "imagenet-val" not in data and "imagenet-v2" not in data:
-        return {}
-    if args.zeroshot_frequency == 0:
-        return {}
-    if (epoch % args.zeroshot_frequency) != 0 and epoch != args.epochs:
-        return {}
-
-    logging.info("Starting zero-shot imagenet.")
-
-    logging.info("Building zero-shot classifier")
-    classifier = zero_shot_classifier(
-        model, imagenet_classnames, openai_imagenet_template, args
-    )
-
-    logging.info("Using classifier")
-    results = {}
-    if "imagenet-val" in data:
-        top1, top5 = run(model, classifier, data["imagenet-val"].dataloader, args)
-        results["imagenet-zeroshot-val-top1"] = top1
-        results["imagenet-zeroshot-val-top5"] = top5
-    if "imagenet-v2" in data:
-        top1, top5 = run(model, classifier, data["imagenet-v2"].dataloader, args)
-        results["imagenetv2-zeroshot-val-top1"] = top1
-        results["imagenetv2-zeroshot-val-top5"] = top5
-
-    logging.info("Finished zero-shot imagenet.")
-
-    return results

models/clap_encoder.py

@@ -1,17 +1,14 @@
-import random
 import torch
 import torch.nn as nn
-import torchaudio
 from models.CLAP.open_clip import create_model
-from models.CLAP.training.data import get_audio_features
 from transformers import RobertaTokenizer
+# removed: import random, torchaudio, get_audio_features — only used by _get_audio_embed (audio modality), not needed for text inference
 
 
 class CLAP_Encoder(nn.Module):
     def __init__(
         self,
         pretrained_path='checkpoint/music_speech_audioset_epoch_15_esc_89.98.pt',
-        sampling_rate=32000,
         amodel = "HTSAT-base",
     ):
         super().__init__()
@@ -22,7 +19,7 @@ class CLAP_Encoder(nn.Module):
         self.enable_fusion = False  # False if you do not want to use the fusion model
         self.fusion_type = "aff_2d"
         self.pretrained = pretrained_path
-        self.sampling_rate = sampling_rate
+        # removed: self.sampling_rate — only used by _get_audio_embed
         self.tokenize = RobertaTokenizer.from_pretrained("roberta-base")
         
         self.model, self.model_cfg = create_model(
@@ -41,39 +38,7 @@ class CLAP_Encoder(nn.Module):
         self.model.eval()
         self.encoder_type = 'CLAP'
 
-    def batch_to_list(self, batch):
-        ret = []
-        for i in range(batch.size(0)):
-            ret.append(batch[i])
-        return ret
-
-    def _get_audio_embed(self, batch):
-        # batch: [B, samples]
-        with torch.no_grad():
-            audio_dict_list = []
-            assert (
-                self.sampling_rate == 32000
-            ), "We only support 32000 sampling rate"
-                
-            # batch: [bs, 1, t-samples]
-            batch = torchaudio.functional.resample(
-                batch, orig_freq=self.sampling_rate, new_freq=48000
-            )
-            for waveform in self.batch_to_list(batch):
-                audio_dict = {}
-                audio_dict = get_audio_features(
-                    audio_dict,
-                    waveform,
-                    480000,
-                    data_truncating="fusion",
-                    data_filling="repeatpad",
-                    audio_cfg=self.model_cfg["audio_cfg"],
-                )
-                audio_dict_list.append(audio_dict)
-                # [bs, 512]
-            embed = self.model.get_audio_embedding(audio_dict_list)
-            
-            return embed.detach()
+    # removed: batch_to_list, _get_audio_embed — audio modality not used in inference
 
     def _get_text_embed(self, batch):
         double_batch = False
@@ -91,15 +56,9 @@ class CLAP_Encoder(nn.Module):
 
 
     def get_query_embed(self, modality, audio=None, text=None, use_text_ratio=0.5, device=None):
-        if modality == 'audio':
-            embed = self._get_audio_embed(audio)
-        elif modality == 'text':
+        # removed: audio and hybird modality branches — only text modality used in inference
+        if modality == 'text':
             embed = self._get_text_embed(text)
-        elif modality == 'hybird':
-            if random.random() > use_text_ratio:
-                embed = self._get_audio_embed(audio)
-            else:
-                embed = self._get_text_embed(text)
         else:
             raise NotImplementedError("Please check flag 'training_modality'.")
 

models/resunet.py

@@ -652,65 +652,6 @@ class ResUNet30(nn.Module):
 
         return output_dict
 
-    @torch.no_grad()
-    def chunk_inference(self, input_dict):
-        chunk_config = {
-                    'NL': 1.0,
-                    'NC': 3.0,
-                    'NR': 1.0,
-                    'RATE': 32000
-                }
-
-        mixtures = input_dict['mixture']
-        conditions = input_dict['condition']
-
-        film_dict = self.film(
-            conditions=conditions,
-        )
-
-        NL = int(chunk_config['NL'] * chunk_config['RATE'])
-        NC = int(chunk_config['NC'] * chunk_config['RATE'])
-        NR = int(chunk_config['NR'] * chunk_config['RATE'])
-
-        L = mixtures.shape[2]
-        
-        out_np = np.zeros([1, L])
-
-        WINDOW = NL + NC + NR
-        current_idx = 0
-
-        while current_idx + WINDOW < L:
-            chunk_in = mixtures[:, :, current_idx:current_idx + WINDOW]
-
-            chunk_out = self.base(
-                mixtures=chunk_in, 
-                film_dict=film_dict,
-            )['waveform']
-            
-            chunk_out_np = chunk_out.squeeze(0).cpu().data.numpy()
-
-            if current_idx == 0:
-                out_np[:, current_idx:current_idx+WINDOW-NR] = \
-                    chunk_out_np[:, :-NR] if NR != 0 else chunk_out_np
-            else:
-                out_np[:, current_idx+NL:current_idx+WINDOW-NR] = \
-                    chunk_out_np[:, NL:-NR] if NR != 0 else chunk_out_np[:, NL:]
-
-            current_idx += NC
-
-            if current_idx < L:
-                chunk_in = mixtures[:, :, current_idx:current_idx + WINDOW]
-                chunk_out = self.base(
-                    mixtures=chunk_in, 
-                    film_dict=film_dict,
-                )['waveform']
-
-                chunk_out_np = chunk_out.squeeze(0).cpu().data.numpy()
-
-                seg_len = chunk_out_np.shape[1]
-                out_np[:, current_idx + NL:current_idx + seg_len] = \
-                    chunk_out_np[:, NL:]
-
-        return out_np
+    # removed: chunk_inference — not called during inference; also had a latent self.sampling_rate AttributeError in the original source