Merge pull request #19 from xinyanghuang7/hxy

添加多图训练和推理过程

.gitignore

@@ -0,0 +1,8 @@
+model/clip_model/clip-vit-base-patch32/
+out/*.pth
+full.json
+trans_json.py
+dataset/*
+!dataset/eval_images/
+!dataset/eval_multi_images/
+minimind-v/model/__pycache__/*

2-sft_vlm.py

@@ -18,7 +18,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from transformers import AutoTokenizer, AutoModel
 from model.model import Transformer
 from model.LMConfig import LMConfig
-from model.dataset import SFTDataset
+from model.dataset import SFTDataset, SFTDataset_multi
 from model.vision_utils import get_vision_model, get_img_embedding
 
 warnings.filterwarnings('ignore')
@@ -97,7 +97,10 @@ def train_epoch(epoch, wandb):
         if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
             model.eval()
             moe_path = '_moe' if lm_config.use_moe else ''
-            ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
+            if args.multi: # 多图训练权重保存
+                ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+            else: # 单图训练权重保存
+                ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
 
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                 state_dict = model.module.state_dict()
@@ -106,12 +109,29 @@ def train_epoch(epoch, wandb):
 
             torch.save(state_dict, ckp)
             model.train()
+    if args.save_last and (not ddp or dist.get_rank() == 0) and args.multi:
+        model.eval()
+        moe_path = '_moe' if lm_config.use_moe else ''
+        ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            state_dict = model.module.state_dict()
+        else:
+            state_dict = model.state_dict()
+
+        torch.save(state_dict, ckp)
+        model.train()
 
 
 def init_model(lm_config):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     moe_path = '_moe' if lm_config.use_moe else ''
-    ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
+
+    # 多图推理建议在单图推理之后
+    if args.multi:
+        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
+    else:
+        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
 
     model = Transformer(lm_config)
     state_dict = torch.load(ckp, map_location=args.device)
@@ -159,6 +179,8 @@ if __name__ == "__main__":
     parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
     parser.add_argument("--data_path", type=str, default="./dataset/LLaVA-Instruct/llava_instruct_230k.json",
                         help="Path to training data")
+    parser.add_argument("--data_path_multi", type=str, default="./dataset/sft_multi_images/output.json",
+                        help="Path to multi images training data")
     parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel")
     parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
     parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
@@ -166,6 +188,8 @@ if __name__ == "__main__":
     parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
     parser.add_argument("--save_interval", type=int, default=100, help="Model saving interval")
     parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training')
+    parser.add_argument('--multi', type=bool, default=False, help='multi-images training')
+    parser.add_argument('--save_last', type=bool, default=True, help='save last step model')
 
     args = parser.parse_args()
 
@@ -197,9 +221,15 @@ if __name__ == "__main__":
     model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
 
     use_version = 0
-    train_ds = SFTDataset(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
+    if args.multi:
+        print("进行多图训练，建议在指令微调后进行...")
+        train_ds = SFTDataset_multi(args.data_path_multi, tokenizer, vision_model=(vision_model, preprocess),
                           image_special_token=lm_config.image_special_token,
                           max_length=max_seq_len)
+    else:
+        train_ds = SFTDataset(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
+                            image_special_token=lm_config.image_special_token,
+                            max_length=max_seq_len)
     train_sampler = DistributedSampler(train_ds) if ddp else None
     train_loader = DataLoader(
         train_ds,

3-eval_chat.py

@@ -1,11 +1,8 @@
 import os
 import random
-import time
-
 import numpy as np
 import torch
 import warnings
-
 from PIL import Image
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from model.model import Transformer
@@ -19,14 +16,16 @@ def count_parameters(model):
     return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
 
-def init_model(lm_config):
+def init_model(lm_config, device, multi):
     tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
     model_from = 1  # 1从权重，2用transformers
 
     if model_from == 1:
         moe_path = '_moe' if lm_config.use_moe else ''
-        ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
+        if multi:
-
+            ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
+        else:
+            ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
         model = Transformer(lm_config)
         state_dict = torch.load(ckp, map_location=device)
 
@@ -35,125 +34,103 @@ def init_model(lm_config):
         for k, v in list(state_dict.items()):
             if k.startswith(unwanted_prefix):
                 state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
-
-        # 加载到模型中
         model.load_state_dict(state_dict, strict=False)
     else:
-        model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small',
+        model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small', trust_remote_code=True)
-                                                     trust_remote_code=True)
 
     model = model.to(device)
-
     print(f'模型参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
 
-    (vision_model, preprocess) = get_vision_model()
+    vision_model, preprocess = get_vision_model()
     vision_model = vision_model.to(device)
-    return model, tokenizer, (vision_model, preprocess)
+    return model, tokenizer, vision_model, preprocess
 
 
 def setup_seed(seed):
-    random.seed(seed)  # 设置 Python 的随机种子
+    random.seed(seed)
-    np.random.seed(seed)  # 设置 NumPy 的随机种子
+    np.random.seed(seed)
-    torch.manual_seed(seed)  # 设置 PyTorch 的随机种子
+    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)  # 为当前 GPU 设置随机种子（如果有）
+    torch.cuda.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)  # 为所有 GPU 设置随机种子（如果有）
+    torch.cuda.manual_seed_all(seed)
-    torch.backends.cudnn.deterministic = True  # 确保每次返回的卷积算法是确定的
+    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False  # 关闭 cuDNN 的自动调优，避免不确定性
+    torch.backends.cudnn.benchmark = False
 
 
 if __name__ == "__main__":
-    # -----------------------------------------------------------------------------
+    # -------------------------- 基本参数设置 -----------------------------------
+    multi = False  # 设置multi参数，控制单图或多图推理
     out_dir = 'out'
-    start = ""
     temperature = 0.5
     top_k = 8
     setup_seed(1337)
-    # device = 'cpu'
     device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
     dtype = 'bfloat16'
-    max_seq_len = 1 * 1024
+    max_seq_len = 1024
     lm_config = LMConfig()
     lm_config.max_seq_len = max_seq_len
-    # -----------------------------------------------------------------------------
+    model, tokenizer, vision_model, preprocess = init_model(lm_config, device, multi)
+    model.eval()
+    
+    # -------------------------- 问题和目录设置 -----------------------------------
+    if multi:
+        image_dir = './dataset/eval_multi_images/bird/'
+        prompt = "<image>\n<image>\nName all the differences between these two birds."
+    else:
+        image_dir = './dataset/eval_images/'
+        prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
 
-    model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
-
-    model = model.eval()
-    # 推送到huggingface
-    # model.push_to_hub("minimind")
-    # tokenizer.push_to_hub("minimind")
-
-    stream = True
-
-    prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
-
-    i = 0
-    # 获取图像文件列表
-    image_dir = './dataset/eval_images/'
     image_files = sorted(os.listdir(image_dir))
 
-    # 遍历图像文件
+    # -------------------------- 推理逻辑 -----------------------------------
-    for image_file in image_files:
+    if multi:
-        image_path = os.path.join(image_dir, image_file)
+        # 多图推理：所有图像编码一次性推理
-
+        image_encoders = []
-        # 打开图像并转换为RGB格式
+        for image_file in image_files:
-        image = Image.open(image_path).convert('RGB')
+            image_path = os.path.join(image_dir, image_file)
-        image_process = get_img_process(image, preprocess).to(vision_model.device)
+            image = Image.open(image_path).convert('RGB')
-        # 对图像进行编码
+            image_process = get_img_process(image, preprocess).to(vision_model.device)
-        image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
+            image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
-
+            image_encoders.append(image_encoder)
-        print(f'[Image]: {image_file}')
+            print(f'[Image]: {image_file}')
-        prompt_ = prompt.replace("\n", "")
+        image_encoders = torch.cat(image_encoders, dim=0).unsqueeze(0)
-        print(f'[Q]: {prompt_}')
 
         messages = [{"role": "user", "content": prompt}]
-
+        new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)[-(max_seq_len - 1):]
-        # print(messages)
-        new_prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )[-(max_seq_len - 1):]
-
         x = tokenizer(new_prompt).data['input_ids']
-        x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
+        x = torch.tensor(x, dtype=torch.long, device=device)[None, ...]
-
-        answer = new_prompt
 
         with torch.no_grad():
             res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
-                                   top_k=top_k, stream=stream, image_encoders=image_encoder)
+                                   top_k=top_k, stream=True, image_encoders=image_encoders)
             print('[A]: ', end='')
-            try:
-                y = next(res_y)
-            except StopIteration:
-                print("No answer")
-                continue
-
             history_idx = 0
-            while y != None:
+            for y in res_y:
                 answer = tokenizer.decode(y[0].tolist())
-                if answer and answer[-1] == '<EFBFBD>':
-                    try:
-                        y = next(res_y)
-                    except:
-                        break
-                    continue
-                # print(answer)
-                if not len(answer):
-                    try:
-                        y = next(res_y)
-                    except:
-                        break
-                    continue
-
                 print(answer[history_idx:], end='', flush=True)
-                try:
-                    y = next(res_y)
-                except:
-                    break
                 history_idx = len(answer)
-                if not stream:
-                    break
-
             print('\n')
+
+    else:
+        # 单图推理：对每张图像单独推理
+        for image_file in image_files:
+            image_path = os.path.join(image_dir, image_file)
+            image = Image.open(image_path).convert('RGB')
+            image_process = get_img_process(image, preprocess).to(vision_model.device)
+            image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
+
+            messages = [{"role": "user", "content": prompt}]
+            new_prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)[-(max_seq_len - 1):]
+            x = tokenizer(new_prompt).data['input_ids']
+            x = torch.tensor(x, dtype=torch.long, device=device)[None, ...]
+
+            print(f'[Image]: {image_file}')
+            with torch.no_grad():
+                res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
+                                       top_k=top_k, stream=True, image_encoders=image_encoder)
+                print('[A]: ', end='')
+                history_idx = 0
+                for y in res_y:
+                    answer = tokenizer.decode(y[0].tolist())
+                    print(answer[history_idx:], end='', flush=True)
+                    history_idx = len(answer)
+                print('\n')

README.md

@@ -180,6 +180,7 @@ CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
       目录下，命名为`*_llm.pth`
     * 3.4 `python 1-pretrain_vlm.py` 执行预训练，得到 `*_vlm_pretrain.pth` 作为预训练的输出权重
     * 3.5 `python 2-sft_vlm.py` 执行指令微调，得到 `*_vlm_sft.pth` 作为指令微调的输出权重
+    * 3.6 `python 2-sft_vlm.py --multi True` 在指令微调的基础上执行多图指令微调（效果仅参考），得到 `*_vlm_sft_multi.pth` 作为多图指令微调的输出权重，512+8 模型显存占用约为8198M
 
 * 4.测试自己训练的模型推理效果
     * 确保需要使用的，训练完成的参数权重`*.pth`文件位于`./out/`目录下
@@ -195,6 +196,7 @@ CPU: Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz
       ```
     * `python 3-eval_chat.py`测试模型的对话效果，其中测试图片在`./dataset/eval_images`下，可自行更换
       ![eval_chat](images/3-eval_chat.png)
+    * `python 3-eval_chat.py`调整[multi](./3-eval_chat.py#L61)变量，可以测试模型的多图对话效果，其中测试图片在`./dataset/eval_multi_images`下，可自行更换（多图数据集规模相对较小且为英文对话，数据集仅包含两图对比的场景，因此微调效果有限）。
 
 🍭 【Tip】预训练和全参指令微调pretrain和sft均支持多卡加速
 
@@ -314,11 +316,76 @@ minimind-v使用50个字符组成的 `<<<...>>>` 占位符代替图像，
 <<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\n这个图片描述的是什么内容？
 ```
 
-计算完embedding和projection，并对图像部分token替换后
+计算完embedding和projection，并对图像部分token替换后整个计算过程到输出则和LLM部分没有任何区别。
-整个计算过程到输出则和LLM部分没有任何区别。
 
 ![input](./images/minimind-v-input.png)
 
+多图实现方法就是通过注入多个\<image\>图像占位符进行实现，不需要修改任何框架。
+
+>  ps: 唯一值得注意的点是，如果在训练过程中存在不同conversations插入图片数量不同的情况，需要利用空特征将较短的特征进行填充（对应[dataset的第267行](./model/dataset.py#L267)），以保证能够在同样大小下被dataloader读取。
+
+> pps: 在prompt中不需要如此做，仍旧是根据插入图像的数量来进行占位符的注入。因此，最终输入给LLM的input feature不会受填充特征的影响。
+
+
+<details>
+<summary> 实现视频理解能力的思考 </summary>
+
+对于多模态大模型的视频理解能力，一个可行的思路是参考现有MiniCPM-V 2.6 进行视频理解的Python示例。
+主要思想是通过提取视频关键帧，而后进行多图推理。
+因此，如果希望在MiniMind-V中添加视频理解能力，可以在现有多图训练的基础上，参考此python脚本中对于关键帧的提取方法，而后加大训练文件中支持图片的数量。
+所支持的MAX_NUM_FRAMES越多，所消耗的显存越大。
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer
+from decord import VideoReader, cpu    # pip install decord
+
+model = AutoModel.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True,
+    attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
+model = model.eval().cuda()
+tokenizer = AutoTokenizer.from_pretrained('openbmb/MiniCPM-V-2_6', trust_remote_code=True)
+
+MAX_NUM_FRAMES=64 # if cuda OOM set a smaller number
+
+def encode_video(video_path):
+    def uniform_sample(l, n):
+        gap = len(l) / n
+        idxs = [int(i * gap + gap / 2) for i in range(n)]
+        return [l[i] for i in idxs]
+
+    vr = VideoReader(video_path, ctx=cpu(0))
+    sample_fps = round(vr.get_avg_fps() / 1)  # FPS
+    frame_idx = [i for i in range(0, len(vr), sample_fps)]
+    if len(frame_idx) > MAX_NUM_FRAMES:
+        frame_idx = uniform_sample(frame_idx, MAX_NUM_FRAMES)
+    frames = vr.get_batch(frame_idx).asnumpy()
+    frames = [Image.fromarray(v.astype('uint8')) for v in frames]
+    print('num frames:', len(frames))
+    return frames
+
+video_path="video_test.mp4"
+frames = encode_video(video_path)
+question = "Describe the video"
+msgs = [
+    {'role': 'user', 'content': frames + [question]}, 
+]
+
+# Set decode params for video
+params = {}
+params["use_image_id"] = False
+params["max_slice_nums"] = 2 # 如果cuda OOM且视频分辨率大于448*448可设为1
+
+answer = model.chat(
+    image=None,
+    msgs=msgs,
+    tokenizer=tokenizer,
+    **params
+)
+print(answer)
+```
+</details>
+
 <u>至此，MiniMind-V的所有细节已经呈现完毕。</u>
 
 <u>MiniMind-V的实现未参考 **任何** 第三方代码，完全基于MiniMind尽可能做最小改动产生，故代码实现和LlaVA等模型必然存在很大区别。
@@ -369,16 +436,62 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
 }
 ```
 
+多图指令微调数据集格式：
+
+```json
+{
+    "id": "000000000000",
+    "image": "0.jpg, 1.jpg",
+    "conversations": [
+      {
+        "from": "human",
+        "value": "context: Source Image: <image> Target Image: <image> Instruction: What is the correct image edit instruction that can transfrom the source image to target image?"
+      },
+      {
+        "from": "gpt",
+        "value": "take the people out of the back in the photo. Remove the two people behind the woman in the white dress and the man in the blue suit. remove people behind the couple in the centre"
+      }
+    ]
+  }
+```
+
+```json
+{
+    "id": "000000013633",
+    "image": "27266.jpg, 27267.jpg",
+    "conversations": [
+      {
+        "from": "human",
+        "value": "<image>\n<image>\nName all the differences between these two birds."
+      },
+      {
+        "from": "gpt",
+        "value": "animal1 is brown with white tuft while animal2 is white with dark brown wings"
+      }
+    ]
+  }
+```
+
+注：
++ 对于指令微调，仅保留了一轮对话，训练单轮对话模型，防止小模型性能被长文本拉低。
++ 多图数据集规模相对较小且为英文对话，数据集仅包含两图对比的场景，因此微调效果有限，这里只提供一种参考思路。
+
+
 注：对于指令微调，仅保留了一轮对话，训练单轮对话模型，防止小模型性能被长文本拉低。
 
+
 最终的数据集下载地址：[百度网盘](https://pan.baidu.com/s/1Nz36OBBvVBGEx-PwIb7ofg?pwd=6666) | [HuggingFace](https://huggingface.co/datasets/jingyaogong/minimind-v_dataset)
 
+多图对话数据集：[HuggingFace](https://hf-mirror.com/datasets/xinyanghuang/minimind-v_multi_image/tree/main)
+
 ## 训练
 
 预训练从595K条数据集中学习图片的通用知识，比如鹿是鹿，狗是狗。
 
 指令微调从230K条真实对话数据集中学习对图片提问的真实问答格式。
 
+多图微调提供两个数据集，图像转换数据集和鸟类对比数据集，长度分别为3.5k和13.6k的真实问答格式。
+
 `1-pretrain_vlm.py` 执行预训练，得到 `*_vlm_pretrain.pth` 作为预训练的输出权重。
 
 `2-sft_vlm.py` 执行指令微调，得到 `*_vlm_sft.pth` 作为指令微调的输出权重。
@@ -412,6 +525,7 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
 
 ### 效果测试
 
+#### 单图对话
 <table>
   <thead>
     <tr>
@@ -496,6 +610,25 @@ MiniMind-V与MiniMind的代码核心改动不超过100行，上手难度低。</
   </tbody>
 </table>
 
+#### 多图对话（效果十分有限）
+
+<table>
+  <thead>
+    <tr>
+      <th>图片1</th>
+      <th>图片2</th>
+      <th>512_sft_multi</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td><img src="./dataset/eval_multi_images/bird/0.jpg" alt="a-bird.png" style="width: 200px;"></td>
+      <td><img src="./dataset/eval_multi_images/bird/1.jpg" alt="a-bird.png" style="width: 200px;"></td>
+      <td>animal1 has a brown and black head with a black and white striped head . animal2 has a black head with a white stripe on its wings .</td>
+    </tr>
+  </tbody>
+</table>
+
 ### 启动推理
 
 ```bash

model/LMConfig.py

@@ -7,15 +7,15 @@ class LMConfig(PretrainedConfig):
 
     def __init__(
             self,
-            dim: int = 512,
+            dim: int = 512, # 768
-            n_layers: int = 8,
+            n_layers: int = 8, # 16
             n_heads: int = 16,
             n_kv_heads: int = 8,
             vocab_size: int = 6400,
             hidden_dim: int = None,
             multiple_of: int = 64,
             norm_eps: float = 1e-5,
-            max_seq_len: int = 200,
+            max_seq_len: int = 512,
             dropout: float = 0.0,
             flash_attn: bool = True,
             image_special_token: str = '<' * 25 + '>' * 25,

model/dataset.py

@@ -178,7 +178,101 @@ class SFTDataset(Dataset):
         image_encoders = get_img_process(image, self.preprocess)
 
         return X_tensor, Y_tensor, loss_mask_tensor, image_encoders
+    
 
+class SFTDataset_multi(Dataset):
+    def __init__(self, json_path, tokenizer, vision_model=None, max_length=1024,
+                 prompt_max_len=512,
+                 answer_max_len=256,
+                 image_special_token='<' * 25 + '>' * 25):
 
-if __name__ == "__main__":
+        super().__init__()
-    pass
+        with open(json_path, 'r', encoding='utf-8') as f:
+            self.data = json.load(f)
+
+        self.max_length = max_length
+        self.prompt_max_len = prompt_max_len
+        self.answer_max_len = answer_max_len
+        self.tokenizer = tokenizer
+        self.vision_model, self.preprocess = vision_model
+        self.padding = 0
+        self.bos_id = self.tokenizer('<s>assistant').data['input_ids']
+        self.dataset_path = './dataset/sft_multi_images/'
+        self.image_special_token = image_special_token
+
+    def __len__(self):
+        return len(self.data)
+
+    def find_sublist_index(self, main_list, sub_list) -> int:
+        last_index = -1
+        for i in range(len(main_list) - len(sub_list) + 1):
+            if main_list[i:i + len(sub_list)] == sub_list:
+                last_index = i
+        return last_index
+
+    def safe_eval(self, s):
+        try:
+            res = eval(s)
+        except Exception as e:
+            return []
+        return res
+
+    def __getitem__(self, index: int):
+        sample = self.data[index]
+        image_names = sample['image'].split(', ')
+        conversation = sample['conversations']
+
+        messages = []
+        for i in range(0, 1):
+            if i + 1 < len(conversation):
+                q = conversation[i]['value'].replace('<image>', self.image_special_token)
+                a = conversation[i + 1]['value']
+
+                if q and a:
+                    messages.append({"role": "user", "content": q})
+                    messages.append({"role": "assistant", "content": a})
+
+        # print(messages) # [{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nAre the two images below that resemble each other described by the same term?  You must choose your answer from the Choice List.  Choice_List:  True, False.'}, {'role': 'assistant', 'content': 'False'}][{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nDo these four pictures fall into the same category?  You must choose your answer from the Choice List.  Choice_List:  True, False.'}, {'role': 'assistant', 'content': 'False'}]
+        new_prompt = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        input_id = self.tokenizer(new_prompt).data['input_ids'][:self.max_length]
+        # print(len(input_id)) # 165 or 263 or 259
+
+        question_length = self.find_sublist_index(input_id, self.bos_id) + len(self.bos_id)
+        padding_len = self.max_length - len(input_id)
+        input_id = input_id + [self.padding] * padding_len
+        mask_len = len(input_id) - question_length - padding_len
+        loss_mask = [0] * question_length + [1] * (mask_len) + [0] * padding_len
+
+        input_id = np.array(input_id)
+
+        X = np.array(input_id[:-1]).astype(np.int64)
+        Y = np.array(input_id[1:]).astype(np.int64)
+        loss_mask = np.array(loss_mask[1:]).astype(np.int64)
+
+        X_tensor = torch.from_numpy(X)
+        Y_tensor = torch.from_numpy(Y)
+        loss_mask_tensor = torch.from_numpy(loss_mask)
+
+        # 读取多张图像
+        image_encoders = []
+        for image_name in image_names:
+            image = Image.open(f'{self.dataset_path}{image_name.strip()}')  # 去掉可能的空格
+            image_encoders.append(get_img_process(image, self.preprocess)['pixel_values'])
+
+        # 确定目标形状
+        max_images = 2  # 根据你的需求设置 一次性最大输入几张图片
+        target_shape = (max_images, 3, 224, 224)
+
+        # 创建填充张量
+        padded_image_encoders = torch.zeros(target_shape, dtype=torch.float32)
+
+        # 填充图像编码
+        for i, img_enc in enumerate(image_encoders):
+            if i < max_images:
+                padded_image_encoders[i] = img_enc.squeeze(0)
+
+        return X_tensor, Y_tensor, loss_mask_tensor, padded_image_encoders

model/model.py

@@ -372,30 +372,39 @@ class Transformer(PreTrainedModel):
         # 查找token中<image>片段的索引，为了替换做准备
         def find_indices(tokens, image_ids):
             image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
-            return [
+            indices = []
-                [i, i + len(image_ids) - 1]
-                for batch_idx in range(tokens.size(0))
-                for i in range(tokens.size(1) - len(image_ids) + 1)
-                if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor)
-            ] or None
 
-        image_indices = find_indices(tokens, self.image_ids)
+            for batch_idx in range(tokens.size(0)):
+                for i in range(tokens.size(1) - len(image_ids) + 1):
+                    if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor):
+                        indices.append([batch_idx, i, i + len(image_ids) - 1])  # 返回batch_idx和开始结束索引
+
+            return indices if indices else None
+
+        image_indices = find_indices(tokens, self.image_ids) # [0, 4, 53], [0, 54, 103], [0, 104, 153], [0, 154, 203] or [1, 4, 53], [1, 54, 103]
 
         # 如果此时有图像编码
         if image_encoders is not None:
-            vision_proj = self.vision_proj(image_encoders)
+                vision_proj = self.vision_proj(image_encoders)
-            if image_indices is not None:
+                vision_proj = vision_proj.unsqueeze(0) if len(vision_proj.shape) == 3 else vision_proj
-                # 创建一个新的张量来存储拼接后的结果
+                if image_indices is not None:
-                new_h = []
+                    # 创建一个新的张量来存储拼接后的结果
-                for i in range(h.size(0)):
+                    new_h = []
-                    before = h[i, :image_indices[i][0], :]
+                    for i in range(h.size(0)):
-                    after = h[i, image_indices[i][1] + 1:, :]
+                        # i即为current_batch_idx索引
-                    # 拼接 before, vision_proj, after
+                        img_idx = 0
-                    new_h_i = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
+                        for batch_idx, start_idx, end_idx in image_indices:
-                    new_h.append(new_h_i)
+                            if batch_idx == i:
-                # 将所有拼接后的结果堆叠起来
+                                # 插入vision_proj特征
-                new_h = torch.stack(new_h, dim=0)
+                                before = h[i][:start_idx, :]
-                return new_h
+                                after = h[i][end_idx + 1:, :]
+                                # 拼接 before, vision_proj, after
+                                h[i] = torch.cat((before, vision_proj[i][img_idx], after), dim=0)[:seqlen]
+                                img_idx += 1
+                        new_h.append(h[i])
+                    # 将所有拼接后的结果重新堆叠起来
+                    new_h = torch.stack(new_h, dim=0)  # torch.Size([32, 511, 512])
+                    return new_h
 
         return h
 
@@ -431,7 +440,8 @@ class Transformer(PreTrainedModel):
             self.last_loss = None
 
         self.OUT.__setitem__('logits', logits)
-        self.OUT.__setitem__('last_loss', self.last_loss)
+        # self.OUT.__setitem__('last_loss', self.last_loss)
+        self.OUT.__setitem__('loss', self.last_loss)
         return self.OUT
 
     @torch.inference_mode()

model/vision_utils.py

@@ -3,6 +3,7 @@ from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 import requests
 import torch
+import transformers
 
 warnings.filterwarnings('ignore')
 
@@ -31,7 +32,10 @@ def get_img_embedding(batch_encoding, vision_model):
         embeddings.append(output.last_hidden_state)
 
     # 从 BatchEncoding 中提取图像张量
-    image_tensor = batch_encoding['pixel_values']  # 假设图像张量在 'pixel_values' 键下
+    if isinstance(batch_encoding, transformers.tokenization_utils_base.BatchEncoding):
+        image_tensor = batch_encoding['pixel_values']
+    else:
+        image_tensor = batch_encoding # torch.Size([32, 4, 3, 224, 224])
 
     # 如果图像张量的形状是5维，则无需添加额外维度
     if len(image_tensor.shape) == 4:
@@ -54,5 +58,5 @@ def get_img_embedding(batch_encoding, vision_model):
         hook.remove()
 
     # 拼接所有特征向量成为一个张量
-    all_embeddings = torch.stack(embeddings, dim=0).squeeze()
+    all_embeddings = torch.stack(embeddings, dim=0).squeeze() # torch.Size([32, 4, 50, 768])
     return all_embeddings