diff --git a/eval_vlm.py b/eval_vlm.py index 631fbb0..b8aa5ec 100644 --- a/eval_vlm.py +++ b/eval_vlm.py @@ -16,18 +16,21 @@ def count_parameters(model): def init_model(lm_config, device): - tokenizer = AutoTokenizer.from_pretrained('./model') + + if args.load == 0: + tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer') moe_path = '_moe' if args.use_moe else '' modes = {0: 'pretrain_vlm', 1: 'sft_vlm', 2: 'sft_vlm_multi'} ckp = f'./{args.out_dir}/{modes[args.model_mode]}_{args.hidden_size}{moe_path}.pth' model = MiniMindVLM(lm_config, vision_model_path="./model/vision_model/clip-vit-base-patch16") state_dict = torch.load(ckp, map_location=device) model.load_state_dict({k: v for k, v in state_dict.items() if 'mask' not in k}, strict=False) + else: - transformers_model_path = 'MiniMind2-Small-V' + transformers_model_path = './MiniMind2-V' tokenizer = AutoTokenizer.from_pretrained(transformers_model_path) - model = AutoModelForCausalLM.from_pretrained(transformers_model_path, trust_remote_code=True) + model = MiniMindVLM.from_pretrained(transformers_model_path, trust_remote_code=True) model.vision_encoder, model.processor = MiniMindVLM.get_vision_model("./model/vision_model/clip-vit-base-patch16") print(f'VLM参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万') @@ -62,7 +65,7 @@ if __name__ == "__main__": # 默认单图推理，设置为2为多图推理 parser.add_argument('--use_multi', default=1, type=int) parser.add_argument('--stream', default=True, type=bool) - parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载") + parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载") parser.add_argument('--model_mode', default=1, type=int, help="0: Pretrain模型，1: SFT模型，2: SFT-多图模型 (beta拓展)") args = parser.parse_args() diff --git a/model/model_vlm.py b/model/model_vlm.py index 72db789..e8d9563 100644 --- a/model/model_vlm.py +++ b/model/model_vlm.py @@ -2,10 +2,13 @@ import os import torch import warnings + +from transformers.generation.utils import GenerateOutput + from .model_minimind import * -from typing import Optional, Tuple, List +from typing import Optional, Tuple, List, Callable from torch import nn -from transformers import CLIPProcessor, CLIPModel +from transformers import CLIPProcessor, CLIPModel, GenerationConfig, LogitsProcessorList, StoppingCriteriaList from typing import List warnings.filterwarnings('ignore') @@ -24,6 +27,7 @@ class VLMConfig(MiniMindConfig): self.image_ids = image_ids super().__init__(**kwargs) + class VisionProj(nn.Module): def __init__(self, ve_hidden_size=768, hidden_size=512): super().__init__() @@ -166,3 +170,137 @@ class MiniMindVLM(MiniMindForCausalLM): self.OUT.__setitem__('aux_loss', aux_loss) self.OUT.__setitem__('past_key_values', presents) return self.OUT + + + ## 以下函数使用gpt生成 + def sample_logits(self, logits: torch.Tensor, temperature: float = 1.0, top_p: float = 1.0) -> torch.Tensor: + if temperature != 1.0: + logits = logits / (temperature + 1e-6) + + probs = F.softmax(logits, dim=-1) + + if top_p < 1.0: + sorted_probs, sorted_indices = torch.sort(probs, descending=True) + cumulative_probs = torch.cumsum(sorted_probs, dim=-1) + + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[:, 1:] = sorted_indices_to_remove[:, :-1].clone() + sorted_indices_to_remove[:, 0] = False + + indices_to_remove = torch.zeros_like(probs, dtype=torch.bool) + indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove) + + probs = probs.masked_fill(indices_to_remove, 0.0) + probs = probs / probs.sum(dim=-1, keepdim=True) + + return torch.multinomial(probs, num_samples=1) + + @torch.inference_mode() + def generate( + self, + input_ids: torch.LongTensor, + pixel_values: Optional[torch.FloatTensor] = None, + max_new_tokens: Optional[int] = 128, + do_sample: bool = False, + temperature: float = 1.0, + top_p: float = 1.0, + num_return_sequences: int = 1, + attention_mask: Optional[torch.Tensor] = None, + pad_token_id: int = 0, + eos_token_id: int = 2, + streamer=None, + use_cache: bool = True, + **kwargs + ) -> Union[GenerateOutput, torch.LongTensor]: + + + + generated = [] + + for batch_idx in range(input_ids.size(0)): + input_ids_i = input_ids[batch_idx][input_ids[batch_idx] != pad_token_id].unsqueeze(0) + pixel_values_i = pixel_values[batch_idx:batch_idx + 1] if pixel_values is not None else None + for _ in range(num_return_sequences): + output = self._stream( + input_ids=input_ids_i, + pixel_values=pixel_values_i, + eos_token_id=eos_token_id, + max_new_tokens=max_new_tokens, + temperature=temperature, + top_p=top_p, + use_cache=use_cache, + do_sample=do_sample, + streamer=streamer, + ) + tokens = [] + for t in output: + if streamer is None: + tokens.append(t[:, -1:]) + if tokens: + result = torch.cat(tokens, dim=1) + full_sequence = torch.cat([input_ids_i, result], dim=1) + else: + full_sequence = input_ids_i + generated.append(full_sequence) + + max_len = max(seq.shape[1] for seq in generated) + generated = [ + torch.cat([ + seq, + torch.full((1, max_len - seq.shape[1]), pad_token_id, dtype=seq.dtype, device=seq.device) + ], dim=1) + for seq in generated + ] + return torch.cat(generated, dim=0) + + def _stream( + self, + input_ids: torch.LongTensor, + pixel_values: Optional[torch.FloatTensor] = None, + eos_token_id: Optional[int] = 2, + max_new_tokens: int = 128, + temperature: float = 1.0, + top_p: float = 1.0, + do_sample: bool = False, + streamer=None, + use_cache: bool = True, + **kwargs + ): + past_key_values = None + first_pass = True + start_pos = input_ids.shape[1] + for _ in range(max_new_tokens): + if first_pass: + out = self( + input_ids=input_ids, + pixel_values=pixel_values, + past_key_values=None, + use_cache=use_cache, + **kwargs + ) + first_pass = False + else: + out = self( + input_ids=input_ids[:, -1:], # only last token + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs + ) + + logits = out.logits[:, -1, :] + past_key_values = out.past_key_values + + if do_sample: + next_token = self.sample_logits(logits) + else: + next_token = torch.argmax(logits, dim=-1, keepdim=True) + + input_ids = torch.cat([input_ids, next_token], dim=1) + + if streamer is not None: + streamer.put(next_token) + + yield input_ids[:, start_pos:] + + if next_token.item() == eos_token_id: + break diff --git a/scripts/web_demo_vlm.py b/scripts/web_demo_vlm.py index 3d0fce2..c564372 100644 --- a/scripts/web_demo_vlm.py +++ b/scripts/web_demo_vlm.py @@ -192,7 +192,7 @@ if __name__ == '__main__': parser.add_argument('--max_seq_len', default=8192, type=int) parser.add_argument('--use_moe', default=False, type=bool) parser.add_argument('--stream', default=True, type=bool) - parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载") + parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载") args = parser.parse_args() lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py index aca682a..c0e6bdb 100644 --- a/trainer/train_pretrain_vlm.py +++ b/trainer/train_pretrain_vlm.py @@ -55,6 +55,7 @@ def train_epoch(epoch, wandb): scaler.scale(loss).backward() + if (step + 1) % args.accumulation_steps == 0: scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip) @@ -98,18 +99,20 @@ def train_epoch(epoch, wandb): def init_model(model_config: VLMConfig): - tokenizer = AutoTokenizer.from_pretrained('../model', use_fast=True) + tokenizer = AutoTokenizer.from_pretrained('../model/minimind_tokenizer', use_fast=True) moe_path = '_moe' if model_config.use_moe else '' # 加载纯语言模型权重 ckp = f'{args.save_dir}/llm_{model_config.hidden_size}{moe_path}.pth' model = MiniMindVLM(model_config, vision_model_path="../model/vision_model/clip-vit-base-patch16") state_dict = torch.load(ckp, map_location=args.device) - model.load_state_dict(state_dict, strict=False) + model.load_state_dict(state_dict, strict=False) # 允许有些参数缺失或多余,即只加载llm部分的参数 - # 冻结除 vision_proj 外的所有参数 - for name, param in model.named_parameters(): - if 'vision_proj' not in name: - param.requires_grad = False + # 冻结除 vision_proj lm_head外的所有参数 + # for name, param in model.named_parameters(): + # if 'vision_proj' in name or 'lm_head' in name: + # param.requires_grad = True + # else: + # param.requires_grad = False Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')