TinyLLaVA工厂

🎉 新闻

• [2025.01] 我们的新工作 TinyLLaVA-Video 发布。
• [2024.08.13] 添加了一个用于解释 TinyLLaVA 预测的简单 可视化工具。
• [2024.05.21] 我们的论文：TinyLLaVA工厂：面向小型多模态模型的模块化代码库 发表！
• [2024.05.15] TinyLLaVA工厂，我们的新代码库，发布！ 请注意，旧代码库 TinyLLaVABench 已移至 tinyllava_bench 分支。
• [2024.05.04] TinyLLaVA演示 发布！（访问我们演示的密码是‘1234’。）
• [2024.02.21] 我们的论文：TinyLLaVA：小型多模态模型框架 发表！

🔥 要点

• 我们最好的模型，TinyLLaVA-Phi-2-SigLIP-3.1B，在整体性能上优于现有的7B模型，如 LLaVA-1.5 和 Qwen-VL。

• TinyLLaVA工厂是一个开源的模块化代码库，专为小型多模态模型（LMMs）设计，基于 PyTorch 和 HuggingFace 实现，注重代码实现的简洁性、新功能的可扩展性以及训练结果的可重复性。

• 使用 TinyLLaVA工厂，您可以以更少的编码工作和更少的错误来定制您自己的大型多模态模型。

• TinyLLaVA工厂集成了多种前沿模型和方法。

  • LLM 目前支持 OpenELM、TinyLlama、StableLM、Qwen、Gemma 和 Phi。

  • 视觉塔目前支持 CLIP、SigLIP、Dino 以及 CLIP 和 Dino 的组合。

  • 连接器目前支持 MLP、Qformer 和 Resampler。

  • 训练配方目前支持 冻结/完全/部分微调 以及 LoRA/QLoRA 微调。

目录

• 🎉 新闻
• 🔥 要点
• 目录
• 安装与要求
  • 升级到最新代码库
• 开始使用
  • 1. 数据准备
  • 2. 训练
  • 3. 评估
• 模型库
  • 已训练模型
    • 模型性能
  • 遗留模型
• 本地启动演示
  • Gradio 网页演示
  • CLI 推理
  • 快速推理脚本
• 自定义微调
• 定制您自己的大型多模态模型
  • LLM
  • 视觉塔
  • 连接器
• 致谢
• 联系方式
• ✏ 引用
• ❤️ 社区努力

安装与要求

请注意，我们的环境要求与 LLaVA 的环境要求不同。我们强烈建议您按照以下步骤从头创建环境。

1. 克隆此仓库并进入文件夹

git clone https://github.com/TinyLLaVA/TinyLLaVA_Factory.git
cd TinyLLaVA_Factory

2. 创建一个 conda 环境，激活它并安装包

conda create -n tinyllava_factory python=3.10 -y
conda activate tinyllava_factory
pip install --upgrade pip  # 启用 PEP 660 支持
pip install -e .

3. 安装额外的包

pip install flash-attn==2.5.7 --no-build-isolation

升级到最新代码库

git pull
pip install -e .

开始使用

1. 数据准备

请参阅我们 文档 中的 数据准备 部分。

2. 训练

这里有一个使用 Phi-2 训练 LMM 的示例。

• 在 scripts/train/train_phi.sh 中将数据路径替换为您自己的路径
• 在 scripts/train/pretrain.sh 中将 output_dir 替换为您自己的路径
• 在 scripts/train/finetune.sh 中将 pretrained_model_path 和 output_dir 替换为您自己的路径
• 在 scripts/train/pretrain.sh 和 scripts/train/finetune.sh 中调整您的 GPU ID（localhost）以及 per_device_train_batch_size

bash scripts/train/train_phi.sh

以下是预训练和微调中使用的重要超参数。

提示：

全局批量大小 = GPU 数量 * per_device_train_batch_size * gradient_accumulation_steps，我们建议您在对模型进行 LoRA 微调之外，始终将全局批量大小和学习率保持在上述水平。

conv_version 是一个用于为不同 LLM 选择不同聊天模板的超参数。在预训练阶段，所有 LLM 的 conv_version 都相同，使用 pretrain。而在微调阶段，我们使用：

phi 用于 Phi-2、StableLM、Qwen-1.5

llama 用于 TinyLlama、OpenELM

gemma 用于 Gemma

3. 评估

请参阅我们 文档 中的 评估 部分。

模型库

训练好的模型

这些模型使用 TinyLLaVA Factory 进行训练。

• TinyLLaVA-Phi-2-SigLIP-3.1B
• TinyLLaVA-Gemma-SigLIP-2.4B
• TinyLLaVA-OpenELM-450M-SigLIP-0.89B
• TinyLLaVA-Qwen2-0.5B-SigLIP
• TinyLLaVA-Qwen2.5-3B-SigLIP

模型性能

旧版模型

这些模型使用旧代码库 TinyLLaVABench 进行训练。

• TinyLLaVA-3.1B
• TinyLLaVA-2.0B
• TinyLLaVA-1.5B
• tiny-llava-hf

如果您有使用我们旧代码库 TinyLLaVABench 训练的模型，并且仍然希望使用它们，我们提供了一个 TinyLLaVA-3.1B 的示例，说明如何使用旧版模型。

from tinyllava.eval.run_tiny_llava import eval_model
from tinyllava.model.convert_legecy_weights_to_tinyllavafactory import *

model = convert_legecy_weights_to_tinyllavafactory('bczhou/TinyLLaVA-3.1B')

prompt = "What are the things I should be cautious about when I visit here?"
image_file = "https://llava-vl.github.io/static/images/view.jpg"

args = type('Args', (), {
    "model_path": None,
    "model": model,
    "query": prompt,
    "conv_mode": "phi", # 与训练阶段的 conv_version 相同。不同的 LLM 有不同的 conv_mode/conv_version，请根据实际情况替换。
    "image_file": image_file,
    "sep": ",",
    "temperature": 0,
    "top_p": None,
    "num_beams": 1,
    "max_new_tokens": 512
})()

eval_model(args)

"""
输出： 
When visiting this serene lakeside location with a wooden dock, there are a few things to be cautious about. First, ensure that the dock is stable and secure before stepping onto it, as it might be slippery or wet, especially if it's a wooden structure. Second, be mindful of the surrounding water, as it can be deep or have hidden obstacles, such as rocks or debris, that could pose a risk. Additionally, be aware of the weather conditions, as sudden changes in weather can make the area more dangerous. Lastly, respect the natural environment and wildlife, and avoid littering or disturbing the ecosystem.
"""

在本地启动演示

Gradio Web 演示

通过运行以下命令启动本地 Web 演示：

python tinyllava/serve/app.py --model-path tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B

CLI 推理

我们还支持使用 CLI 进行推理。要使用我们的模型，请运行：

python -m tinyllava.serve.cli \
   --model-path tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B \
   --image-file "./tinyllava/serve/examples/extreme_ironing.jpg" 

快速推理脚本

如果您想在本地运行自己训练或我们提供的模型，这里有一个示例。

from tinyllava.eval.run_tiny_llava import eval_model

model_path = "/绝对路径/到/您的/模型/"
prompt = "当我访问这里时，有哪些事情需要特别注意？"
image_file = "https://llava-vl.github.io/static/images/view.jpg"
conv_mode = "phi" # 或 llama、gemma 等

args = type('Args', (), {
    "model_path": model_path,
    "model": None,
    "query": prompt,
    "conv_mode": conv_mode,
    "image_file": image_file,
    "sep": ",",
    "temperature": 0,
    "top_p": None,
    "num_beams": 1,
    "max_new_tokens": 512
})()

eval_model(args)

from transformers import AutoTokenizer, AutoModelForCausalLM

hf_path = 'tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B'
model = AutoModelForCausalLM.from_pretrained(hf_path, trust_remote_code=True)
model.cuda()
config = model.config
tokenizer = AutoTokenizer.from_pretrained(hf_path, use_fast=False, model_max_length = config.tokenizer_model_max_length,padding_side = config.tokenizer_padding_side)
prompt="这是什么？"
image_url="http://images.cocodataset.org/val2017/000000039769.jpg"
output_text, genertaion_time = model.chat(prompt=prompt, image=image_url, tokenizer=tokenizer)

print('模型输出:', output_text)
print('运行时间:', genertaion_time)

自定义微调

如果您想使用自己的数据集对 TinyLLaVA 进行微调，请参阅这里。

定制您自己的大型多模态模型

LLM

如果您想自行添加一个新的 LLM，您需要在 tinyllava/data/template/ 和 tinyllava/model/llm/ 文件夹下创建两个文件：一个用于聊天模板，另一个用于语言模型。

以下是以添加 Gemma 模型为例。

首先，创建 tinyllava/data/template/gemma_template.py，该文件将在微调阶段使用。

from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
from packaging import version

from .formatter import EmptyFormatter, StringFormatter
from .base import Template
from .formatter import Formatter
from . import register_template
from ...utils.constants import *

from transformers import PreTrainedTokenizer
import torch
import tokenizers

    
system = "一位好奇的用户与人工智能助手之间的对话。助手会针对用户的问题给出有帮助、详细且礼貌的回答。"

@register_template('gemma') # 使 TemplateFactory 能够通过字符串 ('gemma') 获取所添加的模板。
@dataclass
class GemmaTemplate(Template):
    format_image_token: "Formatter" = StringFormatter(slot="<image>\n{{content}}")
    format_user: "Formatter" = StringFormatter(slot="USER" + ": " + "{{content}}" + " ")
    format_assistant: "Formatter" = StringFormatter(slot="ASSISTANT" + ": " + "{{content}}" + "<eos>") # 根据您选择的分词器进行修改
    system: "Formatter" = EmptyFormatter(slot=system+" ")
    separator: "Formatter" = EmptyFormatter(slot=[' ASSISTANT: ', '<eos>']) # 根据您选择的分词器进行修改

    def _make_masks(self, labels, tokenizer, sep, eos_token_length, rounds):
        # 您的代码在这里
        return labels, cur_len

提示：

请确保 _make_masks 函数返回的 labels 遵循以下格式：答案和 EOS 标记 ID 不被掩码，其他标记则用 -100 掩码。

其次，创建 tinyllava/model/llm/gemma.py。

from transformers import GemmaForCausalLM, AutoTokenizer
# 您想要添加的 LLM 及其对应的分词器。

from . import register_llm

# 添加 GemmaForCausalLM 及其对应的分词器，并处理特殊标记。
@register_llm('gemma') # 使 LLMFactory 能够通过字符串 ('gemma') 获取所添加的 LLM。
def return_gemmaclass(): 
    def tokenizer_and_post_load(tokenizer):
        tokenizer.pad_token = tokenizer.unk_token
        return tokenizer
    return (GemmaForCausalLM, (AutoTokenizer, tokenizer_and_post_load))

最后，创建 scripts/train/train_gemma.sh，并指定相应的 LLM_VERSION 和 CONV_VERSION。

视觉塔

如果您想添加一个新的视觉塔，您需要实现一个新的视觉塔类，该类应继承自基类 VisionTower。以下是 MoF 视觉塔的示例。

首先，创建 tinyllava/model/vision_tower/mof.py：

@register_vision_tower('mof')      
class MoFVisionTower(VisionTower):
    def __init__(self, cfg):
        super().__init__(cfg)

        self._vision_tower = MoF(cfg)
        self._image_processor = # 您的图像处理器
  
    def _load_model(self, vision_tower_name, **kwargs):
        # 您的代码在这里，确保您的模型能够通过 Hugging Face 或 PyTorch 加载方式正确地从预训练参数中加载

    def forward(self, x, **kwargs):
        # 您的代码在这里

然后，根据相应的 VT_VERSION 修改您的训练脚本。

连接器

如果您想添加一个新的连接器，您需要实现一个新的连接器类，该类应继承自基类 Connector。以下是线性连接器的示例。

首先，创建 tinyllava/model/connector/linear.py：

import torch.nn as nn

from . import register_connector
from .base import Connector
    
@register_connector('linear') # 使 ConnectorMFactory 能够通过字符串 ('linear') 获取所添加的连接器。     
class LinearConnector(Connector):
    def __init__(self, config):
        super().__init__()
        self._connector =  nn.Linear(config.vision_hidden_size, config.hidden_size) # 定义您的连接器模型

然后，根据相应的 CN_VERSION 修改您的训练脚本。

致谢

我们特别感谢 Lei Zhao、Luche Wang、Kaijun Luo 和 Junchen Wang 构建了演示。

联系方式

如果您有任何问题，欢迎随时发起 Issue 或通过微信（微信号：TinyLLaVA）联系我们。

✏ 引用

如果您在研究中发现我们的论文和代码很有用，请考虑给项目点个赞 :star: 并引用我们 :pencil:。

@misc{zhou2024tinyllava,
      title={TinyLLaVA：小型多模态大模型框架}, 
      author={周百川、胡颖、翁曦、贾俊龙、罗杰、刘希恩、吴继、黄磊},
      year={2024},
      eprint={2402.14289},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

@article{jia2024tinyllava,
  title={TinyLLaVA 工厂：面向小型多模态大模型的模块化代码库},
  author={贾俊龙、胡颖、翁曦、史一鸣、李淼、张兴建、周百川、刘子宇、罗杰、黄磊、吴继},
  journal={arXiv 预印本 arXiv:2405.11788},
  year={2024}
}

❤️ 社区贡献

• 我们的代码库基于 LLaVA 项目构建。非常出色的工作！
• 我们项目使用了来自 ShareGPT4V 项目的数据。非常出色的工作！

TinyLLaVA_Factory 快速上手指南

TinyLLaVA_Factory 是一个用于构建小规模大型多模态模型（LMM）的模块化开源代码库。它基于 PyTorch 和 HuggingFace，支持多种主流 LLM（如 Phi, Qwen, Gemma, OpenELM 等）和视觉编码器（如 CLIP, SigLIP），旨在以最小的代码量实现模型的自定义训练与推理。

环境准备

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

• 操作系统: Linux (推荐)
• Python 版本: 3.10
• GPU: 支持 CUDA 的 NVIDIA GPU（建议显存 16GB 以上以进行训练）
• 依赖管理: Conda

注意：本项目的环境依赖与原版 LLaVA 不同，强烈建议从头创建独立的虚拟环境。

安装步骤

1. 克隆仓库

git clone https://github.com/TinyLLaVA/TinyLLaVA_Factory.git
cd TinyLLaVA_Factory

2. 创建并激活 Conda 环境

conda create -n tinyllava_factory python=3.10 -y
conda activate tinyllava_factory

3. 安装基础依赖

启用 PEP 660 支持并安装项目包：

pip install --upgrade pip
pip install -e .

4. 安装 Flash Attention

为了获得最佳性能，需安装特定版本的 flash-attn：

pip install flash-attn==2.5.7 --no-build-isolation

国内加速提示：如果下载缓慢，可添加清华或阿里镜像源： pip install flash-attn==2.5.7 --no-build-isolation -i https://pypi.tuna.tsinghua.edu.cn/simple

5. 更新代码（可选）

若需获取最新功能，可执行：

git pull
pip install -e .

基本使用

以下是使用 TinyLLaVA_Factory 进行模型推理的最简流程。

启动本地 Web Demo

您可以快速启动一个基于 Gradio 的网页界面来体验模型（以官方发布的 3.1B 模型为例）：

python tinyllava/serve/app.py --model-path tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B

运行后，终端会显示本地访问地址（通常为 http://localhost:7860），在浏览器打开即可上传图片并进行对话。

命令行推理 (CLI)

如果您希望通过命令行直接获取结果，可以使用内置的评估脚本。以下是一个简单的 Python 调用示例：

from tinyllava.eval.run_tiny_llava import eval_model

# 配置参数
args = type('Args', (), {
    "model_path": "tinyllava/TinyLLaVA-Phi-2-SigLIP-3.1B",
    "model": None,
    "query": "Describe this image in detail.",
    "conv_mode": "phi",  # 根据使用的 LLM 类型调整，如 llama, gemma, qwen 等
    "image_file": "https://llava-vl.github.io/static/images/view.jpg",
    "sep": ",",
    "temperature": 0,
    "top_p": None,
    "num_beams": 1,
    "max_new_tokens": 512
})()

# 执行推理
eval_model(args)

快速训练示例

若要训练自己的模型，需先准备数据（参考官方文档 Data Preparation 部分），然后修改脚本中的路径参数并运行：

1. 编辑 scripts/train/train_phi.sh，替换数据路径和输出目录。
2. 执行训练脚本：

bash scripts/train/train_phi.sh

注：预训练和微调的关键超参数（如 Global Batch Size 和学习率）已在脚本中预设，通常无需更改，除非使用 LoRA 微调。