LLaVA系列③——微调训练LLaVA并进行推理（附详细代码+讲解）

引言：感谢 B站Up主 良睦路程序员 ！这篇博客，主要源于他发布的视频 《训练LLaVA模型（数据集构建、基于Trainer的训练框架搭建）——LLaVA系列》 和我个人的一些总结。⭐️ ⭐️

LLaVA系列的上一篇文章链接：LLaVA系列②——从底层构建LLaVA并测试运行（附详细代码+讲解）

LLaVA系列的下一篇文章链接：LLaVA系列④——如何LLaVA的高级方法（附详细代码+讲解）【暂时还在编辑中】

---

✅ NLP 研 2 选手的学习笔记

笔者简介：Wang Linyong，NPU，2023级，计算机技术
研究方向：文本生成、大语言模型
参考视频：https://www.bilibili.com/video/BV1Si421v7j1
论文链接：https://arxiv.org/abs/2304.08485

---

---

1 写在最前面

● 这篇文章是在 《LLaVA系列①——LLaVA的快速学习和简单调用》 和 《LLaVA系列②——从底层构建LLaVA并测试运行》 这 2 篇博客的基础上，做的进一步延伸，也就是如何自己写代码、使用数据集、构建整体的训练框架，来微调属于我们自己的 LLaVA 模型。

● 如果对于多模态大模型 LLaVA 没有一定的知识基础，建议看一下以上提到的这 2 篇博客，补充必要的背景知识。

● 我使用的深度学习的 硬件环境 如下：

项	版本
CUDA	12.2（通过 “nvcc -V” 查看的）
显卡	2 张 NVIDIA GeForce RTX 4090（总共 48 GB显存）
Ubuntu	22.04（通过 “lsb_release -a” 查看的）
驱动	535.171.04（通过 “nvidia-smi” 查看的）
CPU	Intel® Core™ i9-14900KF（通过 “lscpu” 查看的）

● 我使用的深度学习的部分 重要软件环境 如下：

python			3.10.14
torch			2.2.1
transformers	4.45.0
deepspeed		0.15.4
datasets		2.18.0

提醒： 尽量保持 transformers 和 deepspeed 的版本和我的一致，不然后面可能会报错。

---

2 下载数据集

● 首先，我们需要下载，一会用于微调 LLaVA 的 “图像+文本” 的多模态数据集。

● 下载网址：https://huggingface.co/datasets/liuhaotian/LLaVA-CC3M-Pretrain-595K/tree/main

● 如果下载速度比较慢，建议复制下载的链接地址（右击那个“下载按钮”，会有一个提示“复制链接地址”）到迅雷下载。

● 下载好的数据集传入服务器中，我的目录路径如下。也就是把这 5 个文件放入到 train_llava/LLaVA-CC3M-Pretrain-5955K 中（其他文件和文件夹你可以不用管）。

● 接着使用以下 linux 命令，将 images.zip 中的 59.5万 张图片解压到 images_dl 文件夹中：

unzip images.zip -d images_dl

---

3 数据集的读取

3.1 整体概述

● 这部分代码主要实现了一个数据集类 LlavaDataset（用于处理包含图像和文本问答数据）以及一个数据整理器类 TrainLLavaModelCollator（用于将数据处理成适合模型训练的格式）。

● 该部分代码可以写在名为 data.py 文件中，放置的地方为 train_llava/LLaVA-CC3M-Pretrain-5955K 的同级目录：

3.2 详细步骤

● 第一步 | QA数据类的定义： 定义 QaImageOutput 数据类，用于存储问题文本、图像像素值、答案文本对应的 token ids。

● 第二步 | 数据集构建：

1. 初始化： LlavaDataset类 初始化时接收数据集目录路径，调用 build_dataset() 方法。
2. 读取数据： build_dataset() 读取 chat.json 文件和对应图像目录，将 JSON 数据转为字典列表存储对话信息。
3. 获取样本： __getitem__() 根据索引从对话数据中提取问题、答案和对应图像路径。

● 第三步 | 单样本处理函数： build_qaimage() 函数处理单个样本，构建对话消息模板，生成提示文本，读取图像，用处理器处理文本和图像，将答案文本转换为 token ids，最后返回 QaImageOutput对象。

● 第四步 | 数据整理器处理：

1. 初始化： TrainLLavaModelCollator类 初始化时接收处理器和忽略索引。
2. 单样本转换： convert_one_piece() 拼接问题、答案和 “结束符号的token ids”，得到 “输入的token ids”。接着构建 “标签labes”，其中，问题部分用 “忽略索引(即-100)” 填充。
3. 批量处理： __call__() 处理一个批次(batch)的样本，遍历批次中的每个样本，调用 build_qaimage() 和 convert_one_piece()，记录每个样本的输入长度，填充 “输入的token ids” 和 “标签的token ids” 到统一长度，拼接图像像素值，构建注意力掩码，最后返回处理好的批量数据字典。

● 第五步 | 测试验证：

1. 数据集测试： 创建 LlavaDataset 实例，打印数据集样本总数和特定样本示例。
2. 数据整理器测试： 加载处理器，创建 TrainLLavaModelCollator 实例，构建一个小批次数据，调用数据整理器处理，打印处理后数据的键和内容。

3.3 完整代码

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Tuple
from PIL import Image
from torch.utils.data import Dataset
from transformers import AutoProcessor, LlavaProcessor
import pandas as pd
import torch


# 定义数据类，用于存储问答图像输出的结构
@dataclass
class QaImageOutput:
    q_input_ids: torch.Tensor  # 问题文本对应的输入token ids
    pixel_values: torch.Tensor  # 图像像素值张量
    a_input_ids: torch.Tensor  # 答案文本对应的输入token ids


# 自定义数据集类，继承自PyTorch的Dataset
class LlavaDataset(Dataset):
    def __init__(self, dataset_dir: str) -> None:
        """
            初始化数据集
        Args:
            dataset_dir (str): 数据集目录路径
        """
        super().__init__()

        self.chat_data, self.image_dir = self.build_dataset(dataset_dir)  # 构建数据集，返回对话数据和图像目录

    def build_dataset(self, data_dir: str) -> Tuple[List[Dict], Path]:
        """
        构建数据集

        Args:
            data_dir (str): 数据集目录路径

        Returns:
            Tuple[List[Dict], Path]: 对话数据列表和图像目录路径
        """
        data_path = Path(data_dir)   # 转换为Path对象
        chat_file = data_path.joinpath("chat.json")  # 对话数据文件路径
        image_dir = data_path.joinpath("images_dl")  # 图像目录路径

        # 读取JSON文件并转换为字典列表
        chat_data = pd.read_json(chat_file).to_dict(orient="records")  # [{'id': 'GCC_train_002582585', 'image': 'GCC_train_002582585.jpg', 'conversations': [...]}, ...] 

        return chat_data, image_dir

    def __len__(self):
        """返回数据集样本数量"""
        return len(self.chat_data)

    def __getitem__(self, index) -> Tuple[str, str, Path]:
        """
        获取指定索引的样本

        Args:
            index (int): 样本索引

        Returns:
            Tuple[str, str, Path]: 包含问题文本、答案文本和图像路径的元组
        """
        # 获取当前数据项
        cur_data = self.chat_data[index]  # {'id': 'GCC_train_002109690', 'image': 'GCC_train_002109690.jpg', 'conversations': [{...}, {...}]}
        # 获取对话内容
        conversations = cur_data.get("conversations")  # [{'from': 'human', 'value': 'Offer a succinct explanation of the picture presented.\n<image>'}, {'from': 'gpt', 'value': "it 's appropriate for teens to want to spend more time with their peers than their parents as they get older ."}]

        # 提取人类输入(即问题)【其中 <image> 是图像的占位符】
        human_input = conversations[0].get("value")   # 'Offer a succinct explanation of the picture presented.\n<image>'
        # 提取标准的机器人输出(即答案)
        chatbot_output = conversations[1].get("value")  # "it 's appropriate for teens to want to spend more time with their peers than their parents as they get older ."

        # 构建图像路径
        image_path = self.image_dir.joinpath(cur_data.get("image"))  # PosixPath('train_llava/LLaVA-CC3M-Pretrain-595K/images_dl/GCC_train_002109690.jpg')
        return human_input, chatbot_output, image_path


def build_qaimage(processor: AutoProcessor, q_text: str, a_text: str, image_path: Path):
    """
    构建问答图像输入数据

    Args:
        processor (AutoProcessor): 处理器对象
        q_text (str): 问题文本
        a_text (str): 答案文本
        image_path (Path): 图像路径

    Returns:
        QaImageOutput: 包含输入数据的QaImageOutput对象
    """
    # 构建对话消息模板
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": q_text},
    ]

     # 应用对话模板生成prompt
    prompt = processor.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )  # "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\nOffer a succinct explanation of the picture presented.\n<image><|im_end|>\n<|im_start|>assistant\n"

    raw_image = Image.open(image_path)  # 打开图像文件 <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=224x224 at 0x7BD9094E4BE0>
    inputs = processor(raw_image, prompt, return_tensors="pt")  # 使用处理器处理文本和图像 inputs.keys(): dict_keys(['input_ids', 'attention_mask', 'pixel_values'])

    # 处理答案文本生成输入的 token id
    a_input_ids = processor.tokenizer(
        a_text,
        return_tensors="pt",
        padding="longest",
        truncation=True,
    )["input_ids"]  # tensor([[  275,   364,    82,  8311,   369, 26202,   311,  1366,   311,  8329, 803,   882,   448,   862, 25029,  1091,   862,  6562,   438,   807, 633,  9014,   659]]) 

    # 返回包含所有输入数据的对象
    res = QaImageOutput(
        q_input_ids=inputs.get("input_ids"),
        pixel_values=inputs.get("pixel_values"),
        a_input_ids=a_input_ids,
    )
    return res


class TrainLLavaModelCollator:
    def __init__(self, processor: AutoProcessor, IGNORE_INDEX: int) -> None:
        """
        初始化数据整理器

        Args:
            processor (AutoProcessor): 处理器对象
            IGNORE_INDEX (int): 忽略索引值（通常为-100）
        """
        self.processor = processor
        self.ingnore_index = IGNORE_INDEX

    def convert_one_piece(
        self,
        q_input_ids: torch.Tensor,
        a_input_ids: torch.Tensor,
        # pixel_values: torch.Tensor,
    ):
        """
        转换单个样本为模型输入格式

        Args:
            q_input_ids (torch.Tensor): 问题的 token ids
            a_input_ids (torch.Tensor): 答案的 token ids

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: 拼接后的输入 token ids 和标签 labels
        """
        # 拼接问题、答案和结束的 token ids
        input_ids = torch.concat(
            [
                q_input_ids,  # 其中, 151646 就是 <image> 的 token id → tensor([[151644, 8948, 198, 2610, ..., 624, 151646, 151645, 198, 151644, 77091, 198]])
                a_input_ids,  # tensor([[275, 364, 82, 8311, ..., 9014, 659]])
                torch.tensor(self.processor.tokenizer.eos_token_id).reshape(1, -1),
            ],
            axis=1,
        ) # 得到的结果: tensor([[151644, 8948, 198, 2610, ..., 624, 151646, 151645, 198, 151644, 77091, 198, 275, 364, 82, 8311, ..., 9014,    659, 151645]])

        # 构建标签 labels ：问题部分用 IGNORE_INDEX符号 填充，答案部分保留不变
        labels = torch.concat(
            [
                torch.full(q_input_ids.shape, self.ingnore_index),
                a_input_ids,
                torch.tensor(self.processor.tokenizer.eos_token_id).reshape(1, -1),
            ],
            axis=1,
        ) # 得到的结果: tensor([[-100, -100, ..., -100, -100, -100, 275, 364, 82, 8311, ..., 9014, 659, 151645]])

        return input_ids, labels

    def __call__(self, features: List) -> Dict[str, torch.Tensor]:
        """
        处理批次数据。每次取一个 batch_size 的数据时，会调用这个函数

        Args:
            features (List): 一个列表装着的 batch

        Returns:
            Dict[str, torch.Tensor]: 包含一个 batch_size 数据的字典
        """
        input_ids_list = []
        labels_list = []
        pixel_values = []
        max_input_len_list = []

        for feature in features:
            # 构建单个样本的输入数据
            qaimage_output = build_qaimage(
                self.processor, feature[0], feature[1], feature[2]
            )

            # 转换为模型的输入格式
            temp_input_ids, temp_labels = self.convert_one_piece(
                qaimage_output.q_input_ids, qaimage_output.a_input_ids
            )

            # 记录最大长度
            max_input_len_list.append(temp_input_ids.shape[1])  # 比如: [53, 59]
            # 保存中间结果
            input_ids_list.append(temp_input_ids)
            labels_list.append(temp_labels)
            pixel_values.append(qaimage_output.pixel_values)

        # 获取每个 batch 的最大长度
        max_input_len = max(max_input_len_list)

         # 填充输入 token ids 到统一长度
        final_input_ids = torch.concat(
            [
                torch.concat(
                    [
                        torch.full(
                            (1, max_input_len - max_input_len_list[index]),
                            self.processor.tokenizer.pad_token_id,
                        ),
                        value,
                    ],
                    axis=1,
                )
                for index, value in enumerate(input_ids_list)
            ]
        )  # input_ids_list: [tensor([[151644, 8948, ..., 659, 151645]]), tensor([[151644, 8948, 198, 2610, 525, ..., 25956, 151645]])] → final_input_ids: [tensor([[151643, 151643, ..., 151643, 151644, 8948, ..., 659, 151645]]), tensor([[151644, 8948, 198, 2610, 525, ..., 25956, 151645]])]

        # 填充 标签labels 的 token ids 到统一长度
        final_labels = torch.concat(
            [
                torch.concat(
                    [
                        torch.full(
                            (1, max_input_len - max_input_len_list[index]),
                            self.ingnore_index,
                        ),
                        value,
                    ],
                    axis=1,
                )
                for index, value in enumerate(labels_list)
            ]
        )  # 同理获取 final_input_ids 的流程, 不过padding符号变为了-100, 而 final_input_ids 的padding符号为 151643

        # 拼接图像的像素值
        final_pixel_values = torch.concat(pixel_values, axis=0)
        # 构建注意力掩码
        attention_mask = torch.ones_like(final_input_ids)  # tensor([[1, 1, 1, ... 1, 1, 1], [1, 1, 1, 1, 1, 1, ..., 1, 1, 1]])
        attention_mask[final_input_ids == self.processor.tokenizer.pad_token_id] = 0  # tensor([[0, 0, ..., 0, 1, 1, ... 1, 1, 1], [1, 1, ..., 1, 1, 1]])

        return {
            "input_ids": final_input_ids,
            "labels": final_labels,
            "pixel_values": final_pixel_values,
            "attention_mask": attention_mask,
        }


if __name__ == "__main__":
    # 1. 测试 LlavaDataset 类
    data_dir = "train_llava/LLaVA-CC3M-Pretrain-595K"  # 数据集路径（注意：当前路径可能需要调整）
    llavadataset = LlavaDataset(data_dir)   # 创建数据集实例
    all_data_len = len(llavadataset)   # 打印数据集总样本数
    print("all_data_len:", all_data_len)  # all_data_len: 595375
    one_data_example = llavadataset[168]  # 打印第168个样本示例 
    print("one_data_example:", one_data_example)  # one_data_example: ('Offer a succinct explanation of the picture presented.\n<image>', "it 's appropriate for teens to want to spend more time with their peers than their parents as they get older .", PosixPath('train_llava/LLaVA-CC3M-Pretrain-595K/images_dl/GCC_train_002109690.jpg'))

    # 2. 测试 TrainLLavaModelCollator 类
    model_name_or_path = "./my_llava_model/model_01"  # 定义模型的路径，这里指定了本地存储的模型目录
    llava_processor = LlavaProcessor.from_pretrained(model_name_or_path)  # 处理器会加载模型对应的分词器和图像处理器等信息
    tlmc = TrainLLavaModelCollator(processor=llava_processor, IGNORE_INDEX=-100)
    # 自行构建一个 batch_size = 2 的批次
    one_origin_batch = [llavadataset[168], llavadataset[178]]
    # 得到的 one_input_batch 将会传给模型进行训练
    one_input_batch = tlmc(one_origin_batch)  
    print("one_input_batch.keys():", one_input_batch.keys())  # 看一下存在哪些键 → dict_keys(['input_ids', 'labels', 'pixel_values', 'attention_mask'])
    print("one_input_batch:", one_input_batch)
    

3.4 运行结果

● 提醒： 数据集、模型的文件路径需要自己定义好。

---

4 辅助工具函数

● 这部分代码实现了统计并打印 模型可训练参数数量、所有参数数量及可训练参数百分比 的功能。在主程序中加载 LLaVA 模型并应用 LoRA 配置后，调用该函数可以输出该模型的参数统计信息。

● 该部分代码可以写在名为 util.py 文件中，放置的地方为 train_llava/LLaVA-CC3M-Pretrain-5955K 的同级目录：

● 完整代码如下：

import torch.nn as nn

# 代码复制自：https://github.com/huggingface/peft/blob/2f5360a7da22a236b5ad4c059572fff5321c867c/src/peft/peft_model.py#L617
def get_nb_trainable_parameters(model:nn.Module) -> tuple[int, int]:
    """
    返回模型中可训练参数的数量和所有参数的数量。

    参数:
        model (nn.Module): 要统计参数的模型

    返回:
        tuple[int, int]: 一个元组，包含可训练参数的数量和所有参数的数量
    """
    # 初始化可训练参数的数量
    trainable_params = 0
    # 初始化所有参数的数量
    all_param = 0
    # 遍历模型的所有命名参数
    for _, param in model.named_parameters():
        # 获取当前参数的元素数量
        num_params = param.numel()
        # 如果使用 DeepSpeed Zero 3 并且权重初始化为空
        if num_params == 0 and hasattr(param, "ds_numel"):
            # 使用 DeepSpeed 统计的元素数量
            num_params = param.ds_numel

        # 由于 bitsandbytes 库中 4 位线性层的设计
        # 需要将参数数量乘以 2 以获得正确的参数数量
        if param.__class__.__name__ == "Params4bit":
            if hasattr(param, "element_size"):
                # 获取元素的字节大小
                num_bytes = param.element_size()
            elif not hasattr(param, "quant_storage"):
                # 如果没有量化存储属性，默认字节大小为 1
                num_bytes = 1
            else:
                # 获取量化存储的字节大小
                num_bytes = param.quant_storage.itemsize
            # 调整参数数量
            num_params = num_params * 2 * num_bytes

        # 累加所有参数的数量
        all_param += num_params
        # 如果当前参数需要梯度更新
        if param.requires_grad:
            # 累加可训练参数的数量
            trainable_params += num_params

    return trainable_params, all_param


# 代码复制自：https://github.com/huggingface/peft/blob/2f5360a7da22a236b5ad4c059572fff5321c867c/src/peft/peft_model.py#L647
def print_trainable_parameters(model: nn.Module) -> None:
    """
    打印模型中可训练参数的数量。

    注意：print_trainable_parameters() 使用了 get_nb_trainable_parameters()，这与 huggingface/transformers 中的
    num_parameters(only_trainable=True) 不同。get_nb_trainable_parameters() 返回的是 Peft 模型的
    （可训练参数，所有参数），其中包括修改后的骨干变压器模型。对于像 LoRA 这样的技术，骨干变压器模型会被 LoRA 模块就地修改。
    然而，对于提示调优，骨干变压器模型不会被修改。num_parameters(only_trainable=True) 返回的是骨干变压器模型的
    可训练参数数量，这可能会有所不同。

    参数:
        model (nn.Module): 要打印可训练参数信息的模型
    """
    # 获取可训练参数和所有参数的数量
    trainable_params, all_param = get_nb_trainable_parameters(model)

    # 打印可训练参数数量、所有参数数量以及可训练参数的百分比
    print(
        f"可训练参数: {trainable_params:,d} || 所有参数: {all_param:,d} || 可训练百分比: {100 * trainable_params / all_param:.4f}"
    )


if __name__ == "__main__":
    from transformers import LlavaForConditionalGeneration
    from peft import LoraConfig, get_peft_model
    import torch

    # 加载LLaVA模型
    model_name_or_path = "../my_llava_model/model_01"
    model = LlavaForConditionalGeneration.from_pretrained(
        pretrained_model_name_or_path=model_name_or_path,
        torch_dtype=torch.bfloat16,  # 使用BF16精度
        low_cpu_mem_usage=True,  # 优化CPU内存使用
        local_files_only=True   # 仅使用本地文件
    )

    # LoRA参数配置
    LORA_R = 32  # 秩参数，控制低秩近似的维度
    # LORA_ALPHA = 16  # 缩放因子，用于调整 LoRA 模块中权重更新的幅度
    LORA_DROPOUT = 0.05  # dropout率
    TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]  # 需要应用LoRA的模块名称

    # 初始化LoRA配置
    config = LoraConfig(
        r=LORA_R,
        # lora_alpha=LORA_ALPHA,
        target_modules=TARGET_MODULES,
        lora_dropout=LORA_DROPOUT,
        bias="none",  # LoRA 只作用于权重矩阵，而不影响偏置项
        task_type="CAUSAL_LM",  # 表示任务类型是因果语言模型（Causal Language Modeling），即根据前面的文本预测下一个单词
        modules_to_save=["multi_modal_projector"],  # 表示在保存模型时，除了 LoRA 相关的参数外，还需要保存名称为 multi_modal_projector 的模块的参数
    )

    # 应用LoRA到模型
    model = get_peft_model(model, config)

    # 打印可训练参数信息
    print_trainable_parameters(model)

● 运行结果如下，其中训练的参数占整个模型的百分比为 0.9322%：

---

5 模型的训练

5.1 整体概述

● 这部分代码实现了基于 LLaVA 模型的训练功能，支持多种训练方式（全量参数训练、使用 LoRA 训练、冻结视觉塔训练），通过命令行参数配置模型、数据及训练相关设置，完成数据加载、模型训练并保存训练结果。

● 这段代码可以写在名为 simple_LLaVA_run.py 文件中，放置的地方为 train_llava 的同级目录：

5.2 详细步骤

● 第一步 | 导入库和模块： 导入日志记录、数据类定义、PyTorch、Hugging Face 的transformers库等，还导入自定义的数据处理模块和打印可训练参数的工具函数。

● 第二步 | 定义参数配置类：

1. ModelArguments：定义模型相关参数，包括模型路径和训练类型。
2. DataArguments：定义数据相关参数，即训练数据路径。

● 第三步 | 加载模型和处理器函数：

1. 根据 ModelArguments 加载 LLaVA 模型和处理器。
2. 根据不同训练类型进行相应配置。

● 第四步 | 加载数据集和数据整理器函数：

1. 根据 DataArguments 加载自定义的 LlavaDataset 数据集。
2. 初始化 TrainLLavaModelCollator 数据整理器，用于处理数据格式。

● 第五步 | 训练函数：

1. 使用 HfArgumentParser 解析命令行参数。
2. 调用上述函数加载模型、处理器、数据集和数据整理器。
3. 初始化 Trainer 对象，配置模型、训练参数、数据集和数据整理器。
4. 执行训练，保存训练状态和模型到指定输出目录。

● 第六步 | 主程序入口： 配置日志格式，调用训练函数启动训练流程。

5.3 完整代码

import logging
from dataclasses import dataclass, field
from typing import Optional
import torch
import transformers
from transformers import (
    LlavaForConditionalGeneration,
    LlavaProcessor,
    Trainer,
    TrainingArguments,
)
# 导入自定义的数据处理模块（需确保路径正确）
from train_llava.data import LlavaDataset, TrainLLavaModelCollator
from train_llava.util import print_trainable_parameters

logger = logging.getLogger(__name__)


@dataclass
class ModelArguments:
    """模型参数配置类"""
    model_name_or_path: Optional[str] = field(default="test_model/model001")
    train_type: Optional[str] = field(
        default="none",
        metadata={
            "help": """
            1. use_lora: 使用lora训练,
            2. none: 全量参数训练;
            3. freeze_vision: 只冻结vision_tower进行训练
            """
        },
    )


@dataclass
class DataArguments:
    """数据参数配置类"""
    data_path: str = field(
        default=None, metadata={"help": "训练数据的路径"}
    )


def load_model_processor(modelargs: ModelArguments):
    """
    加载模型和处理器
    :param modelargs: 模型参数配置
    :return: 模型和处理器对象
    """
    # 加载LLaVA模型
    model = LlavaForConditionalGeneration.from_pretrained(
        modelargs.model_name_or_path,
        torch_dtype=torch.bfloat16,  # 使用BF16精度
        low_cpu_mem_usage=True,  # 优化CPU内存使用
        local_files_only=True   # 仅使用本地文件
    )
    print("modelargs.model_name_or_path:", modelargs.model_name_or_path)
    # 加载模型对应的处理器（包含分词器和图像处理器）
    processor = LlavaProcessor.from_pretrained(modelargs.model_name_or_path, local_files_only=True)

    if modelargs.train_type == "use_lora":
        logging.warning("Loading model to Lora")

        from peft import LoraConfig, get_peft_model

        # LoRA参数配置
        LORA_R = 32  # 秩参数，控制低秩近似的维度
        # LORA_ALPHA = 16  # 缩放因子，用于调整 LoRA 模块中权重更新的幅度
        LORA_DROPOUT = 0.05  # dropout率
        TARGET_MODULES = ["q_proj", "k_proj", "v_proj", "o_proj"]  # 需要应用LoRA的模块名称

        # 初始化LoRA配置
        config = LoraConfig(
            r=LORA_R,
            # lora_alpha=LORA_ALPHA,
            target_modules=TARGET_MODULES,
            lora_dropout=LORA_DROPOUT,
            bias="none",  # LoRA 只作用于权重矩阵，而不影响偏置项
            task_type="CAUSAL_LM",  # 表示任务类型是因果语言模型（Causal Language Modeling），即根据前面的文本预测下一个单词
            modules_to_save=["multi_modal_projector"],  # 表示在保存模型时，除了 LoRA 相关的参数外，还需要保存名称为 multi_modal_projector 的模块的参数
        )

        # 应用LoRA到模型
        model = get_peft_model(model, config)

    elif modelargs.train_type == "none":
        """全量参数训练"""
        logging.warning("使用全量参数进行训练")

        pass
    elif modelargs.train_type == "freeze_vision":
        """冻结视觉塔的所有参数"""
        logging.warning("冻结vision_tower网络层，剩下的网络权重进行训练")
        for param in model.vision_tower.parameters():
            param.requires_grad = False

    # 打印可训练参数信息
    print_trainable_parameters(model)

    return model, processor


def load_dataset_collator(processor, dataargs: DataArguments):
    """
    加载数据集和数据整理器
    :param processor: 模型处理器
    :param dataargs: 数据参数配置
    :return: 数据集和数据整理器对象
    """
    llava_dataset = LlavaDataset(
        dataargs.data_path  # xxxxx/LLaVA-CC3M-Pretrain-595K
    )

    logger.info(f"Loaded dataset from {dataargs.data_path}")

    # 初始化数据整理器
    data_collator = TrainLLavaModelCollator(
        processor=processor,
        IGNORE_INDEX=-100  # 损失计算时忽略的索引值
    )

    return llava_dataset, data_collator


def train():
    """训练模型的主函数"""
    # 初始化参数解析器
    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    # 解析命令行参数
    model_args, data_args, training_args = parser.parse_args_into_dataclasses()  
    # 加载模型和处理器
    model, processor = load_model_processor(model_args)
    # 加载数据集和数据整理器
    train_dataset, data_collator = load_dataset_collator(processor, data_args)
    # 初始化训练器
    trainer = Trainer(
        model=model,
        args=training_args,  # 训练参数
        train_dataset=train_dataset,
        eval_dataset=None,   # 暂时不使用验证集
        data_collator=data_collator,  # 传入数据整理器
    )

    # 开始训练
    logger.info("开始训练...")
    trainer.train()
    logger.info("训练完成")

     # 保存训练状态和模型
    trainer.save_state()
    trainer.save_model(output_dir=training_args.output_dir)
    logger.info(f"模型已保存到 {training_args.output_dir}")


if __name__ == "__main__":
    # 配置日志格式
    logging.basicConfig(
        format="%(asctime)s %(levelname)s [%(name)s] %(message)s",
        level=logging.INFO,
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    # 执行模型训练
    train()

5.4 运行脚本和其他配置文件

● 首先创建一个 ds_zero2_no_offload.json 和 run_zero2.sh 文件，对应的创建位置如下：

● 其中，ds_zero2_no_offload.json 用于存放 “DeepSpeed 框架的配置的” 各种参数：

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 100,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1e-10
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 1e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 1e8,
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

● 各个参数解释如下：

1. fp16 配置：

• • enabled：“auto” 表示自动决定是否启用混合精度训练（使用半精度浮点数 fp16）。如果硬件支持（如 NVIDIA GPU 具有 Tensor Cores）且框架允许，将启用 fp16 训练以提高训练速度和减少内存使用。
• • loss_scale：设置损失缩放因子，用于在混合精度训练中调整损失值的大小，以避免梯度下溢问题。值为 0 时通常表示使用动态损失缩放。
• • loss_scale_window：动态损失缩放时，用于计算损失缩放因子的窗口大小。每 loss_scale_window 次迭代，框架会检查是否需要调整损失缩放因子。
• • initial_scale_power：动态损失缩放的初始缩放因子的幂。例如，16 表示初始缩放因子为 2¹⁶。
• • hysteresis：在动态损失缩放中，用于避免缩放因子频繁调整的滞后值。只有当梯度的变化超过这个滞后值时，才会调整损失缩放因子。
• • min_loss_scale：允许的最小损失缩放因子，防止缩放因子过小导致梯度信息丢失。
    
    

2. zero_optimization 配置：

• • stage：设置为 2 表示使用 ZeRO 优化的第二阶段。在这个阶段，优化器状态（如梯度和动量）会在多个进程之间进行分片，减少每个进程的内存占用。
• • allgather_partitions：设置为 true 时，在优化器状态更新后，会使用全规约（allgather）操作来收集所有进程的分片数据，确保每个进程都有完整的优化器状态信息。
• • allgather_bucket_size：全规约操作时的桶大小，用于控制数据传输的批量大小，以优化内存和通信效率。
• • overlap_comm：设置为 true 时，允许在计算和通信操作之间进行重叠，以提高训练效率。
• • reduce_scatter：设置为 true 时，在反向传播过程中，会使用规约散射（reduce_scatter）操作来合并梯度，减少通信量。
• • reduce_bucket_size：规约散射操作时的桶大小，与 allgather_bucket_size 类似，用于控制数据传输的批量大小。
• • contiguous_gradients：设置为 true 时，确保梯度在内存中是连续的，这有助于提高计算效率。
    
    

3. gradient_accumulation_steps： 设置为 “auto” 表示自动根据训练配置确定梯度累积的步数。梯度累积可以模拟更大的批次大小，减少内存占用，但会增加反向传播的计算时间。
4. gradient_clipping： 设置为 “auto” 表示自动根据训练情况进行梯度裁剪。梯度裁剪用于防止梯度爆炸问题，确保训练的稳定性。
5. steps_per_print： 设置每 2000 步打印一次训练信息，如损失值、梯度等，方便监控训练进度。
6. train_batch_size： 设置为 “auto” 表示自动根据硬件资源和训练配置确定训练批次大小。批次大小会影响训练的稳定性和收敛速度。
7. train_micro_batch_size_per_gpu： 设置为 “auto” 表示自动确定每个 GPU 上的微批次大小。在梯度累积和分布式训练中，微批次大小是一个重要参数，它控制每次反向传播计算梯度的样本数量。
8. wall_clock_breakdown： 设置为 false 表示不启用训练时间的详细分解记录。如果设置为 true，DeepSpeed 会记录训练过程中不同阶段（如前向传播、反向传播、优化器更新等）的时间，用于性能分析。

● 另外，run_zero2.sh 文件的代码如下：

# 因为使用 RTX 4000 系列显卡时，不支持通过 P2P 或 IB 实现更快的通信宽带，需要设置以下两个环境变量
# 禁用 NCCL 的 P2P 通信，以避免可能出现的兼容性问题
export NCCL_P2P_DISABLE="1"
# 禁用 NCCL 的 IB 通信，以适应 RTX 4000 系列显卡的特性
export NCCL_IB_DISABLE="1"

# 设置 Hugging Face 模型仓库的镜像地址，方便下载模型等资源
export HF_ENDPOINT=https://hf-mirror.com

# 使用 deepspeed 工具运行 simple_LLaVA_run.py 脚本
# --include localhost:0,1 表示指定在本地的 0 号和 1 号 GPU 上运行任务
# 注：localhost 代表本地机器，0 和 1 是 GPU 的编号
deepspeed --include localhost:0,1 simple_LLaVA_run.py \
    --deepspeed ds_zero2_no_offload.json \
    --model_name_or_path my_llava_model/model_01 \
    --train_type use_lora \
    --data_path train_llava/LLaVA-CC3M-Pretrain-595K \
    --remove_unused_columns false \
    --bf16 true \
    --fp16 false \
    --dataloader_pin_memory True \
    --dataloader_num_workers 10 \
    --dataloader_persistent_workers True \
    --output_dir output_model_user_lora_simple_train \
    --num_train_epochs 10 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --evaluation_strategy "no" \
    --save_strategy "epoch" \
    --save_total_limit 3 \
    --report_to "tensorboard" \
    --learning_rate 4e-4 \
    --logging_steps 10

● 部分参数的解释如下：

1. --deepspeed ds_zero2_no_offload.json：指定 deepspeed 的配置文件为 ds_zero2_no_offload.json，用于配置分布式训练的参数等。
2. --model_name_or_path my_llava_model/model_01：指定要使用的预训练模型的名称或路径，这里是本地的 my_llava_model/model_01 路径下的模型。
3. --train_type use_lora：设置训练类型为使用 LoRA 进行训练。
4. --data_path train_llava/LLaVA-CC3M-Pretrain-595K：指定训练数据的路径，这里是 train_llava/LLaVA-CC3M-Pretrain-595K 目录下的数据。
5. --remove_unused_columns false：设置为不删除未使用的列，可能是在数据处理时保留所有列。
6. --bf16 true：启用 bf16（Brain Floating Point 16）混合精度训练。
7. --fp16 false：禁用 fp16（Half Precision Floating Point 16）混合精度训练。
8. --dataloader_pin_memory True：设置数据加载器将数据固定在内存中，以提高数据读取速度。
9. --dataloader_num_workers 10：设置数据加载器使用的工作进程数为 10，用于并行加载数据。
10. --dataloader_persistent_workers True：设置数据加载器的工作进程为持久化，避免重复创建进程的开销。
11. --output_dir output_model_user_lora_simple_train：设置训练结果（模型等）的输出目录为 output_model_user_lora_simple_train。
12. --num_train_epochs 10：设置训练的总 epoch 数为 10，即数据将被训练 10 次。
13. --per_device_train_batch_size 1：设置每个 GPU 设备上的训练批次大小为 1。(两张4090也才能支持这个batch_size😇…)
14. --per_device_eval_batch_size 1：设置每个 GPU 设备上的评估批次大小为 1。
15. --gradient_accumulation_steps 8：设置梯度累积的步数为 8，即每 8 个批次累积一次梯度再更新模型参数。
16. --evaluation_strategy "no"：设置评估策略为不进行评估（“no”），即训练过程中不进行验证集评估。
17. --save_strategy "epoch"：设置模型保存策略为每个 epoch 结束时保存模型。
18. --save_total_limit 3：设置最多保存 3 个模型检查点，避免占用过多磁盘空间。
19. --report_to "tensorboard"：设置将训练过程中的指标等信息报告到 TensorBoard 中，方便可视化监控。
20. --learning_rate 4e-4：设置训练的学习率为 4e-4（0.0004），控制模型参数更新的步长。
21. --logging_steps 10：设置每 10 步记录一次训练日志，方便查看训练进度和指标。

● 提醒： 可能需要使用 chmod 命令来为该脚本添加执行权限。

chmod +x run_zero2.sh

● 执行代码：

./run_zero2.sh

5.5 运行结果

● 提醒： 第一次加载模型会很慢，我用了 10 多分钟。另外，全量数据集的训练需要花费 149 个小时，需要的时间太长了。后续我会提取一些轻量的数据集来训练。

---

6 轻量数据集的微调训练

● 我们从 images_dl 文件夹中随机选择 500 张图片，并将它们复制到同一级目录的 mini_500_images 文件夹中，linux命令如下：

find images_dl -type f -name "*.jpg" -o -name "*.jpeg" -o -name "*.png" | shuf -n 500 | xargs -I {} cp {} mini_500_images

● 运行结果：

● 然后，我们需要构建和 mini_500_images(图像) 对应的文本数据 mini_500_chat.json。我们可以使用 new_data_create.py 脚本来完成这个任务，代码如下：

import json
import os

# 读取chat.json文件
with open('chat.json', 'r') as file:
    chat_data = json.load(file)

# 获取mini_500_images文件夹中所有图片的名称
image_folder ='mini_500_images'
image_names = [os.path.basename(f) for f in os.listdir(image_folder) if os.path.isfile(os.path.join(image_folder, f))]

# 提取匹配的数据
matching_data = [item for item in chat_data if item["image"] in image_names]
print("len(matching_data):", len(matching_data))

# 将匹配的数据写入新的mini_500_chat.json文件
with open('mini_500_chat.json', 'w') as file:
    json.dump(matching_data, file, indent=4)

print("数据分配成功！")

● 同时，我们还需要修改 data.py 的一点代码：

● 最后，我们再次执行 run_zero2.sh 就可以。大概需要个 5~6 分钟。

---

7 加载训练好的模型来推理

● 我用的是 20000 条样例来训练的。

● 完整的推理代码 simple_LLaVA_infer.py：

from transformers import LlavaForConditionalGeneration, AutoProcessor
from peft import PeftModel
from PIL import Image
from train_llava.data import LlavaDataset  # 导入自定义数据集类
import torch
import re


# 配置模型路径
raw_model_name_or_path = "my_llava_model/model_01"  # 原始预训练模型路径
peft_model_name_or_path = "output_model_user_lora_simple_train/checkpoint-12500"  # PEFT微调后的模型路径

# 加载原始模型
model = LlavaForConditionalGeneration.from_pretrained(
    raw_model_name_or_path,
    device_map="cuda:0",  # 指定模型加载到第一个GPU
    torch_dtype=torch.bfloat16  # 使用BF16精度
)

# 加载PEFT适配器
model = PeftModel.from_pretrained(
    model,  # 原始模型
    peft_model_name_or_path,  # PEFT模型路径
    adapter_name="peft_v1"  # 适配器名称
)

# 加载处理器（包含分词器和图像处理器）
processor = AutoProcessor.from_pretrained(raw_model_name_or_path)

# 设置模型为评估模式
model.eval()

# 初始化数据集
llavadataset = LlavaDataset("train_llava/LLaVA-CC3M-Pretrain-595K")

# 验证数据集加载
print(len(llavadataset), llavadataset[10])  # 打印数据集长度和第10个样本

# 选择测试样本
testdata = llavadataset[1688]  # 选择第200个样本
print("testdata:", testdata)  # 打印样本内容：(问题文本, 答案文本, 图像路径)

# 可视化图像（需确保路径正确）
Image.open(testdata[2])  # 打开并显示测试样本的图像


def build_model_input(model, processor, testdata: tuple):
    """
    构建模型输入并生成回答

    Args:
        model: 加载的模型
        processor: 模型处理器
        testdata: 包含(问题文本, 答案文本, 图像路径)的元组

    Returns:
        str: 生成的原始回答文本
    """
    # 构建对话模板
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": testdata[0]},  # 使用问题文本
    ]
    
    # 生成prompt
    prompt = processor.tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )

    # 加载图像
    image = Image.open(testdata[2])
    
    # 预处理输入数据
    inputs = processor(text=prompt, images=image, return_tensors="pt")
    
    # 将输入数据移动到模型所在设备
    for tk in inputs.keys():
        inputs[tk] = inputs[tk].to(model.device)
    
    # 生成回答
    generate_ids = model.generate(**inputs, max_new_tokens=100)  # 最大生成50个token

    # 分离生成部分
    generate_ids = [
        oid[len(iids):] for oid, iids in zip(generate_ids, inputs.input_ids)
    ]

    # 解码生成结果
    gen_text = processor.batch_decode(
        generate_ids,
        skip_special_tokens=False,  # 保留特殊标记
        clean_up_tokenization_spaces=True
    )[0]
    
    return gen_text


def process_string(input_str):
    """
    后处理生成的回答文本

    Args:
        input_str: 原始生成文本

    Returns:
        str: 处理后的文本
    """
    # 1. 分割连续两个以上空格的块
    blocks = re.split(r'\s{2,}', input_str.strip())
    
    # 2. 处理每个块
    processed_blocks = []
    for block in blocks:
        # 去除块内空格
        processed_block = block.replace(' ', '')
        # 替换特殊标记
        if processed_block == '<|im_end|>':
            processed_blocks.append('.')
        else:
            processed_blocks.append(processed_block)
    
    # 3. 重新拼接并处理标点
    intermediate_str = ' '.join(processed_blocks)
    processed_str = re.sub(r'\s+([,.])', r'\1 ', intermediate_str)
    
    # 4. 清理多余空格
    processed_str = re.sub(r'\s+', ' ', processed_str).strip()
    processed_str = processed_str.replace("<|im_end|>", ".")  # 彻底移除特殊标记
    
    return processed_str


# 执行推理
res = build_model_input(model, processor, testdata)
print("res:", res)  # 打印原始生成结果

# 后处理结果
clear_res = process_string(res)
print("clear_res:", clear_res)  # 打印处理后的结果

# 合并并保存模型
model = model.merge_and_unload()  # 合并PEFT参数到基础模型
model.save_pretrained("output_model_lora_merge_001")  # 保存合并后的模型
processor.save_pretrained("output_model_lora_merge_001")  # 保存处理器
print("模型保存成功")

● 关键注释说明：

1. 模型加载：

• • 使用device_map指定 GPU 设备
• • BF16 精度用于加速计算和节省内存
• • PEFT 模型加载后需要合并才能保存完整模型

2. 数据处理：

• • 数据集包含图像和文本对
• • build_model_input函数构建对话模板并生成回答
• • 特殊标记<|im_end|>用于标记图像结束位置

3. 后处理逻辑：

• • 分步骤处理生成文本中的空格和特殊标记
• • 使用正则表达式进行字符串清洗
• • 最终结果去除所有特殊标记并优化格式

4. 模型保存：

• • merge_and_unload将 PEFT 参数合并到基础模型
• • 保存完整的模型和处理器以便后续使用

● 提醒：推理时，我们应该把 data.py 中的这两行代码改成这样：

● 推理所输入的图片如下：

● 运行结果： （输出的结果还是蛮准确的）

a tree with blossoming flowers in the spring.
翻译：一颗春天开花的树。

---

8 参考资料

[1] 《训练LLaVA模型（数据集构建、基于Trainer的训练框架搭建）——LLaVA系列》，感谢B站Up主：良睦路程序员

---

9 补充说明

● 若有写得不对的地方，或有疑问，欢迎评论交流。

● LLaVA系列的上一篇文章链接：LLaVA系列②——从底层构建LLaVA并测试运行（附详细代码+讲解）

● LLaVA系列的下一篇文章链接：LLaVA系列④——使用LLaVA的高级方法（附详细代码+讲解）【暂时还在编辑中】

---

⭐️ ⭐️ 写于2025年3月31日 16:29 教研室工位 💻