From f74124f6b91ebd26b344d567766ba19aac83e613 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Sat, 4 Apr 2026 12:43:13 +0800 Subject: [PATCH] [fix] bugs --- model/model_vlm.py | 19 ++++++++++++++++--- trainer/trainer_utils.py | 2 +- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/model/model_vlm.py b/model/model_vlm.py index 696ccd8..eb5e0d9 100644 --- a/model/model_vlm.py +++ b/model/model_vlm.py @@ -113,8 +113,12 @@ class MiniMindVLM(MiniMindForCausalLM): if pixel_values is not None and start_pos == 0: if hasattr(pixel_values, 'keys'): - img_emb = MiniMindVLM.get_image_embeddings(pixel_values, self.vision_encoder) - vision_tensors = self.vision_proj(img_emb) + sample_val = next(iter(pixel_values.values())) + if sample_val.ndim == 5: + bs, num = sample_val.shape[:2] + vision_tensors = self.vision_proj(MiniMindVLM.get_image_embeddings({k: v.flatten(0, 1) for k, v in pixel_values.items()}, self.vision_encoder)).view(bs, num, self.config.image_token_len, -1) + else: + vision_tensors = self.vision_proj(MiniMindVLM.get_image_embeddings(pixel_values, self.vision_encoder)) else: if len(pixel_values.shape) == 6: pixel_values = pixel_values.squeeze(2) @@ -152,4 +156,13 @@ class MiniMindVLM(MiniMindForCausalLM): loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100) output = MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=presents, hidden_states=hidden_states) - return output \ No newline at end of file + return output + + def generate(self, *args, num_return_sequences=1, **kwargs): + if num_return_sequences > 1 and 'pixel_values' in kwargs: + pv = kwargs['pixel_values'] + if hasattr(pv, 'keys'): + kwargs['pixel_values'] = {k: v.repeat(num_return_sequences, *([1] * (v.ndim - 1))) for k, v in pv.items()} + else: + kwargs['pixel_values'] = pv.repeat(num_return_sequences, *([1] * (pv.ndim - 1))) + return super().generate(*args, num_return_sequences=num_return_sequences, **kwargs) \ No newline at end of file diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py index 6e8b72c..7cc9ac3 100644 --- a/trainer/trainer_utils.py +++ b/trainer/trainer_utils.py @@ -81,7 +81,7 @@ def init_vlm_model(vlm_config, from_weight='pretrain_vlm', tokenizer_path='../mo # 2、判断策略 if freeze_llm == 0: for name, param in model.named_parameters(): - if 'vision_model' not in name: + if 'vision_encoder' not in name: param.requires_grad = True elif freeze_llm == 1: for name, param in model.model.named_parameters():