PyTorch实现的Audio Flamingo

Zhifeng Kong, Arushi Goel, Rohan Badlani, Wei Ping, Rafael Valle, Bryan Catanzaro

[演示网站] [演示视频] [ICML海报]

本仓库包含了Audio Flamingo: 一种具有少样本学习和对话能力的新型音频语言模型（ICML 2024）的PyTorch实现。Audio Flamingo是一种新型音频理解语言模型，具有：

强大的音频理解能力，
通过上下文学习和检索快速适应未见任务的能力，以及
强大的多轮对话能力。

我们引入了一系列训练技术、架构设计和数据策略来增强我们模型的这些能力。在各种音频理解任务上的广泛评估证实了我们方法的有效性，创造了新的最先进基准。

代码结构

foundation/文件夹包含基础模型的训练代码。
chat/文件夹包含聊天模型的训练代码，可以执行多轮对话。
inference/文件夹包含基础模型和聊天模型的推理代码。

每个文件夹内的结构主要基于Open Flamingo仓库（提交a05dcba）。每个文件夹都是自包含的，我们预计这些文件夹之间没有交叉依赖。

准备工作

从官方仓库下载Laion-CLAP的源代码。将文件夹重命名为my_laion_clap/并复制到foundation/、chat/、inference/下。将预训练检查点下载到YOUR_DATA_ROOT_DIR/audio-flamingo-data/laion-clap-pretrained/laion_clap/。
从官方仓库下载Microsoft-CLAP的源代码。将文件夹重命名为my_ms_clap/并复制到foundation/、chat/、inference/下。在每个文件夹中，用clap_modified_code/CLAPWrapper.py替换my_ms_clap/msclap/CLAPWrapper.py，这添加了一些处理函数并修复了clapcap的一些bug。将预训练检查点下载到YOUR_DATA_ROOT_DIR/audio-flamingo-data/clap/。
从原始来源下载原始训练和评估数据集。参考foundation/data/README.md和chat/data/README.md获取准备数据的具体说明。

运行代码

我们参考foundation/README.md、chat/README.md和inference/README.md获取训练基础模型、训练聊天模型和进行推理的具体说明，因为它们需要不同的设置。我们使用8个A100 GPU来训练我们的模型。

检查点

checkpoints/文件夹包含基础模型和聊天模型的检查点。
每个模型约17GB。由于git lfs的限制，我们将每个模型分成5个部分。下载后，进入checkpoints/并运行python checkpoint_utils.py合并各个部分。
或者，模型检查点也在HuggingFace上（下载更容易）：https://huggingface.co/nvidia/audio-flamingo。可以git clone此项目或使用huggingface_hub.hf_hub_download函数下载：checkpoint_path = hf_hub_download(repo_id="nvidia/audio-flamingo", filename="foundation(or chat).pt")。
如果您想用这些检查点进行推理，请记得修改inference/configs/*.yaml和inference/inference_examples.py中的绝对路径，以正确加载模型检查点和数据（参见inference/README.md）。
基础模型使用foundation/configs/foundation_pretrain.yaml进行预训练，然后使用foundation/configs/foundation_sft_8_shot.yaml进行微调。
聊天模型使用foundation/configs/foundation_pretrain.yaml进行预训练，然后使用foundation/configs/foundation_sft_4_shot.yaml进行微调，最后使用chat/configs/chat.yaml进行微调。

下游应用

我们使用Audio Flamingo作为合成字幕的数据标注机器。有关合成数据集和许可说明的详细信息，请参见labeling_machine/。

参考文献

每个文件夹（foundation/、chat/、inference/）中的主要训练和推理代码，包括train/、src/、data/和configs/，均修改自Open Flamingo（提交a05dcba）（MIT许可证），该项目借鉴了flamingo-pytorch（MIT许可证）、flamingo-mini（MIT许可证）和open_clip（MIT许可证）。src/helpers.py还包括基于attention-is-all-you-need-pytorch（MIT许可证）的自注意力实现，该项目借鉴了OpenNMT-py（MIT许可证）。我们的代码还依赖于LAION-AI/CLAP（CC0-1.0许可证）和microsoft/CLAP（MIT许可证）。在chat/data/prepare_each_dataset.py中，过滤关键词基于LLARK论文（CC-BY-4.0许可证）和LTU论文（CC-BY-4.0许可证）。

许可证

本仓库中的代码使用MIT许可证（见LICENSE）。
本仓库中的检查点（checkpoints/*.pt）仅供非商业用途。它们受OPT-IML许可证、OpenAI生成数据的使用条款以及每个训练数据集附带的原始许可证的约束。

引用

@article{kong2024audio,
  title={Audio Flamingo: A Novel Audio Language Model with Few-Shot Learning and Dialogue Abilities},
  author={Kong, Zhifeng and Goel, Arushi and Badlani, Rohan and Ping, Wei and Valle, Rafael and Catanzaro, Bryan},
  journal={arXiv preprint arXiv:2402.01831},
  year={2024}
}

audio-flamingo