预训练改为全参数训练，补全模型缺少的generate
Jinshijiming
parent
623a130c34
commit
1ed44bba5d
11
eval_vlm.py
11
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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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} 百万')
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue