OneLLM：一个将所有模态与语言对齐的统一框架

[项目页面] [论文] [Hugging Face 演示🤗] [ModelScope 演示🤖] [模型🤗] [数据]

新闻

• 2024.02.27 OneLLM 被 CVPR 2024 接收！🎉
• 2023.12.01 发布模型权重和推理代码。

目录

• 安装
• 模型
• 演示
• 数据
• 评估
• 训练

安装

1. 将仓库克隆到本地文件夹。

git clone https://github.com/csuhan/OneLLM

cd OneLLM

2. 安装依赖包。

conda create -n onellm python=3.9 -y
conda activate onellm

pip install -r requirements.txt

# 安装 PointNet
cd model/lib/pointnet2
python setup.py install

3. 安装 Apex。（可选）

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

模型

我们在 Hugging Face 上提供了一个预览模型：csuhan/OneLLM-7B。

演示

Hugging Face 演示： csuhan/OneLLM。

本地演示： 假设你已将权重下载到 ${WEIGHTS_DIR}，则运行以下命令以在本地启动 Gradio 演示。

python demos/multi_turn_mm.py --gpu_ids 0 --tokenizer_path config/llama2/tokenizer.model --llama_config config/llama2/7B.json --pretrained_path ${WEIGHTS_DIR}/consolidated.00-of-01.pth

命令行演示：

python demos/cli.py --image_path ${IMAGE_PATH} --gpu_ids 0 --tokenizer_path config/llama2/tokenizer.model --llama_config config/llama2/7B.json --pretrained_path ${WEIGHTS_DIR}/consolidated.00-of-01.pth

数据

更多详情请参阅 Data.md。

评估

更多详情请参阅 Evaluation.md。

训练

图像-文本预训练

单节点 8 卡训练：exps/image_text_pretrain_8gpu.sh

torchrun --nproc_per_node=8 main_pretrain.py \
--epochs 1 --dataset image \
--batch_size 40 --accum_iter 16 \
--model_parallel_size 1 \
--data_parallel sdp \
--save_consolidated \
--llama_type onellm \
--llama_ckpt_dir ${LLAMA_7B_PATH} \
--llama_config config/llama2/7B.json \
--tokenizer_path config/llama2/tokenizer.model \
--auto_resume \
--weight_decay 0.1 --output_dir ${OUTPUT_DIR} \
--warmup_iters 2000 --lr_decay_iters 200000 --lr 5e-5 --min_lr 5e-6 --clip_grad 2 \
--save_freq 1000 \
2>&1 | tee -a ${OUTPUT_DIR}/output.log

多节点 DDP 训练：

同时在 N 个节点上运行 N 个脚本，即可启动多节点 DDP 训练。以下是一个节点的示例脚本：

MASTER_ADDR=IP_ADDRESS_OF_NODE_1
NNODES=N
MASTER_PORT=29500
NPROC_PER_NODE=8

RANK=0 或 1 或 ... 或 N

bash
torchrun \
--nnodes=$NNODES \
--nproc_per_node=8 \
--node_rank=$RANK \
--master_port=$MASTER_PORT \
--master_addr=$MASTER_ADDR \
main_pretrain.py \
--epochs 1 --dataset image \
--batch_size 40 --accum_iter 16 \
--model_parallel_size 1 \
--data_parallel sdp \
--save_consolidated \
--llama_type onellm \
--llama_ckpt_dir ${LLAMA_7B_PATH} \
--llama_config config/llama2/7B.json \
--tokenizer_path config/llama2/tokenizer.model \
--auto_resume \
--weight_decay 0.1 --output_dir ${OUTPUT_DIR} \
--warmup_iters 2000 --lr_decay_iters 200000 --lr 5e-5 --min_lr 5e-6 --clip_grad 2 \
--save_freq 1000 \
2>&1 | tee -a ${OUTPUT_DIR}/output.log

多节点 SLURM 训练：exps/image_text_pretrain_slurm.sh

#!/bin/bash
#SBATCH --gres=gpu:8
#SBATCH -n 16
#SBATCH -N 2
#SBATCH --cpus-per-task=16

srun python -u main_pretrain.py \
--epochs 1 --dataset image \
--batch_size 40 --accum_iter 8 \
--model_parallel_size 1 \
--data_parallel sdp \
--save_consolidated \
--llama_type onellm \
--llama_ckpt_dir ${LLAMA_7B_PATH} \
--llama_config config/llama2/7B.json \
--tokenizer_path config/llama2/tokenizer.model \
--auto_resume \
--weight_decay 0.1 --output_dir ${OUTPUT_DIR} \
--warmup_iters 2000 --lr_decay_iters 200000 --lr 5e-5 --min_lr 5e-6 --clip_grad 2 \
--save_freq 1000 \
2>&1 | tee -a ${OUTPUT_DIR}/output.log

多模态-文本预训练

阶段 II 预训练：假设我们已有预训练好的 ${IMAGE_TEXT_MODEL}，则运行 exps/multimodal_text_pretrain_stage2.sh 进行视频-音频-点云-文本的预训练。

阶段 III 预训练：假设我们已有预训练好的 ${STAGE2_MODEL}，则运行 exps/multimodal_text_pretrain_stage3.sh 进行深度-法线-IMU-fMRI-文本的预训练。

指令微调

假设我们已有预训练好的 ${STAGE3_MODEL}，则运行 exps/multimodal_text_finetune.sh 进行多模态指令微调。

引用

@InProceedings{han2023onellm,
  title={OneLLM: One Framework to Align All Modalities with Language},
  author={Han, Jiaming and Gong, Kaixiong and Zhang, Yiyuan and Wang, Jiaqi and Zhang, Kaipeng and Lin, Dahua and Qiao, Yu and Gao, Peng and Yue, Xiangyu},
  booktitle = {IEEE/CVF 计算机视觉与模式识别会议（CVPR）论文集},
  year={2024}
}

致谢

LLaMA, LLaMA-Adapter, LLaMA2-Accessory, Meta-Transformer, ChatBridge

许可证

本项目基于 Llama 2 开发，请参阅 LLAMA 2 社区许可证。

OneLLM 快速上手指南

OneLLM 是一个统一的框架，旨在将图像、视频、音频、点云、深度图、IMU、fMRI 等多种模态与语言模型对齐。本项目基于 LLaMA 2 构建，并已在 CVPR 2024 上发表。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.9
• GPU: 支持 CUDA 的 NVIDIA GPU (建议显存 16GB 以上用于推理，多卡用于训练)
• 依赖管理: Conda (推荐)
• 网络环境: 能够访问 Hugging Face 或 ModelScope (国内用户建议使用 ModelScope 或配置镜像加速)

安装步骤

1. 克隆代码库

git clone https://github.com/csuhan/OneLLM
cd OneLLM

2. 创建虚拟环境并安装基础依赖

conda create -n onellm python=3.9 -y
conda activate onellm

# 安装 requirements.txt 中的依赖
# 国内用户若下载缓慢，可添加清华源：-i https://pypi.tuna.tsinghua.edu.cn/simple
pip install -r requirements.txt

# 安装 PointNet++ 组件
cd model/lib/pointnet2
python setup.py install
cd ../../..

3. 安装 Apex (可选，推荐用于训练加速)

如果您计划进行模型训练，建议安装 NVIDIA Apex 以启用混合精度训练和扩展算子。

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
cd ..

4. 获取模型权重

您可以从 Hugging Face 或 ModelScope 下载预训练权重。

• Hugging Face: csuhan/OneLLM-7B
• ModelScope (国内加速): 请访问 ModelScope 搜索 OneLLM 相关仓库下载。

假设您将权重下载到了本地目录 ${WEIGHTS_DIR}，且文件名包含 consolidated.00-of-01.pth。

基本使用

安装完成后，您可以通过命令行 (CLI) 或本地 Gradio 界面体验 OneLLM 的多模态对话能力。

方式一：命令行交互 (CLI Demo)

这是最轻量级的使用方式，适合测试单张图片的理解能力。请将 ${IMAGE_PATH} 替换为您本地的图片路径，${WEIGHTS_DIR} 替换为权重所在目录。

python demos/cli.py \
--image_path ${IMAGE_PATH} \
--gpu_ids 0 \
--tokenizer_path config/llama2/tokenizer.model \
--llama_config config/llama2/7B.json \
--pretrained_path ${WEIGHTS_DIR}/consolidated.00-of-01.pth

运行后，您可以在终端输入问题与模型进行多轮对话。

方式二：本地 Web 界面 (Local Gradio Demo)

如果您更喜欢图形化界面，可以启动本地的 Gradio 服务。

python demos/multi_turn_mm.py \
--gpu_ids 0 \
--tokenizer_path config/llama2/tokenizer.model \
--llama_config config/llama2/7B.json \
--pretrained_path ${WEIGHTS_DIR}/consolidated.00-of-01.pth

运行成功后，终端会显示一个本地 URL（通常是 http://127.0.0.1:7860），在浏览器中打开即可上传图像、音频等文件并与模型交互。

注意: 首次运行时，程序会自动下载 LLaMA 2 的分词器配置文件。如果网络受阻，请手动下载 tokenizer.model 和 7B.json 配置到 config/llama2/ 目录下。