250427
jingyaogong
parent
26de815253
commit
623a130c34
11
README.md
11
README.md
|
|
@ -64,7 +64,16 @@
|
|||
### 👉**最近更新**
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-20 (newest 🎉)</b> </summary>
|
||||
<summary> <b>2025-04-27 (newest 🎉)</b> </summary>
|
||||
|
||||
- 兼容性更新
|
||||
- 适配[「minimind仓库新特性」](https://github.com/jingyaogong/minimind/issues/370)
|
||||
- 规范化部分代码
|
||||
|
||||
</details>
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-20</b> </summary>
|
||||
|
||||
- MiniMind2-V伴随MiniMind2同步更新
|
||||
- 大幅减少所有冗余代码，规范代码格式
|
||||
|
|
|
|||
13
README_en.md
13
README_en.md
|
|
@ -7,7 +7,6 @@
|
|||
|
||||
<div align="center">
|
||||
|
||||
|
||||
[](https://github.com/jingyaogong/minimind-v/stargazers)
|
||||
[](LICENSE)
|
||||
[](https://github.com/jingyaogong/minimind-v/commits/master)
|
||||
|
|
@ -71,7 +70,17 @@ Is the training process difficult? Now, let's explore the answers and feel the j
|
|||
### 👉**Recent Updates**
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-20 (newest 🎉)</b> </summary>
|
||||
<summary> <b>2025-04-27 (newest 🎉)</b> </summary>
|
||||
|
||||
- Compatibility updates
|
||||
- Adapted to the new feature in the "minimind" repository
|
||||
- Standardized parts of the code
|
||||
|
||||
</details>
|
||||
|
||||
|
||||
<details close>
|
||||
<summary> <b>2025-02-20</b> </summary>
|
||||
|
||||
- MiniMind2-V updated alongside MiniMind2
|
||||
- Significant reduction of all redundant code, standardized code format
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import json
|
|||
from PIL import Image
|
||||
from torch.utils.data import Dataset, DataLoader
|
||||
import torch
|
||||
from .model_vlm import MiniMindVLM
|
||||
from model.model_vlm import MiniMindVLM
|
||||
import os
|
||||
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
|
|
@ -20,8 +20,8 @@ class VLMDataset(Dataset):
|
|||
self.max_length = max_length
|
||||
self.preprocess = preprocess
|
||||
self.image_token = image_special_token
|
||||
self.bos_id = tokenizer('<s>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>', add_special_tokens=False).input_ids
|
||||
self.bos_id = tokenizer('<|im_start|>assistant', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('<|im_end|>', add_special_tokens=False).input_ids
|
||||
|
||||
def __len__(self):
|
||||
return len(self.samples)
|
||||
88
eval_vlm.py
88
eval_vlm.py
|
|
@ -5,12 +5,8 @@ import numpy as np
|
|||
import torch
|
||||
import warnings
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model_vlm import MiniMindVLM
|
||||
from model.VLMConfig import VLMConfig
|
||||
from transformers import logging as hf_logging
|
||||
|
||||
hf_logging.set_verbosity_error()
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -20,22 +16,23 @@ def count_parameters(model):
|
|||
|
||||
|
||||
def init_model(lm_config, device):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model')
|
||||
if args.load == 0:
|
||||
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.dim}{moe_path}.pth'
|
||||
model = MiniMindVLM(lm_config)
|
||||
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-V'
|
||||
transformers_model_path = 'MiniMind2-Small-V'
|
||||
tokenizer = AutoTokenizer.from_pretrained(transformers_model_path)
|
||||
model = AutoModelForCausalLM.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} 百万')
|
||||
|
||||
vision_model, preprocess = MiniMindVLM.get_vision_model()
|
||||
vision_model, preprocess = model.vision_encoder, model.processor
|
||||
return model.eval().to(device), tokenizer, vision_model.eval().to(device), preprocess
|
||||
|
||||
|
||||
|
|
@ -56,62 +53,59 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--temperature', default=0.65, type=float)
|
||||
parser.add_argument('--top_p', default=0.85, type=float)
|
||||
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
|
||||
# MiniMind2-Small (26M)：(dim=512, n_layers=8)
|
||||
# MiniMind2 (104M)：(dim=768, n_layers=16)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
# MiniMind2-Small (26M)：(hidden_size=512, num_hidden_layers=8)
|
||||
# MiniMind2 (104M)：(hidden_size=768, num_hidden_layers=16)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=8192, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
# 默认单图推理，设置为2为多图推理
|
||||
parser.add_argument('--use_multi', default=1, type=int)
|
||||
parser.add_argument('--stream', default=True, type=bool)
|
||||
parser.add_argument('--load', default=0, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
parser.add_argument('--load', default=1, 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()
|
||||
|
||||
lm_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||
|
||||
lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||
model, tokenizer, vision_model, preprocess = init_model(lm_config, args.device)
|
||||
streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
|
||||
|
||||
def chat_with_vlm(prompt, pixel_tensors, image_names):
|
||||
def chat_with_vlm(prompt, pixel_values, image_names):
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
new_prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-args.max_seq_len + 1:]
|
||||
|
||||
inputs = tokenizer(
|
||||
new_prompt,
|
||||
return_tensors="pt",
|
||||
truncation=True
|
||||
).to(args.device)
|
||||
|
||||
print(f'[Image]: {image_names}')
|
||||
with torch.no_grad():
|
||||
x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
|
||||
outputs = model.generate(
|
||||
x,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
max_new_tokens=args.max_seq_len,
|
||||
temperature=args.temperature,
|
||||
top_p=args.top_p,
|
||||
stream=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
pixel_tensors=pixel_tensors
|
||||
)
|
||||
print('🤖️: ', end='')
|
||||
try:
|
||||
if not args.stream:
|
||||
print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
|
||||
else:
|
||||
history_idx = 0
|
||||
for y in outputs:
|
||||
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
|
||||
if (answer and answer[-1] == '<EFBFBD>') or not answer:
|
||||
continue
|
||||
print(answer[history_idx:], end='', flush=True)
|
||||
history_idx = len(answer)
|
||||
except StopIteration:
|
||||
print("No answer")
|
||||
print('\n')
|
||||
print('🤖️: ', end='')
|
||||
generated_ids = model.generate(
|
||||
inputs["input_ids"],
|
||||
max_new_tokens=args.max_seq_len,
|
||||
num_return_sequences=1,
|
||||
do_sample=True,
|
||||
attention_mask=inputs["attention_mask"],
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
streamer=streamer,
|
||||
top_p=args.top_p,
|
||||
temperature=args.temperature,
|
||||
pixel_values=pixel_values
|
||||
)
|
||||
|
||||
response = tokenizer.decode(generated_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||||
messages.append({"role": "assistant", "content": response})
|
||||
print('\n\n')
|
||||
|
||||
|
||||
# 单图推理：每1个图像单独推理
|
||||
|
|
|
|||
|
|
@ -1,61 +0,0 @@
|
|||
from transformers import PretrainedConfig
|
||||
from typing import List
|
||||
|
||||
|
||||
class LMConfig(PretrainedConfig):
|
||||
model_type = "minimind"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int = 512,
|
||||
n_layers: int = 8,
|
||||
n_heads: int = 8,
|
||||
n_kv_heads: int = 2,
|
||||
vocab_size: int = 6400,
|
||||
hidden_dim: int = None,
|
||||
multiple_of: int = 64,
|
||||
norm_eps: float = 1e-5,
|
||||
max_seq_len: int = 8192,
|
||||
rope_theta: int = 1e6,
|
||||
dropout: float = 0.0,
|
||||
flash_attn: bool = True,
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
use_moe: bool = False,
|
||||
####################################################
|
||||
num_experts_per_tok: int = 2,
|
||||
n_routed_experts: int = 4,
|
||||
n_shared_experts: bool = True,
|
||||
scoring_func: str = 'softmax',
|
||||
aux_loss_alpha: float = 0.1,
|
||||
seq_aux: bool = True,
|
||||
norm_topk_prob: bool = True,
|
||||
**kwargs,
|
||||
):
|
||||
self.dim = dim
|
||||
self.n_layers = n_layers
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.vocab_size = vocab_size
|
||||
self.hidden_dim = hidden_dim
|
||||
self.multiple_of = multiple_of
|
||||
self.norm_eps = norm_eps
|
||||
self.max_seq_len = max_seq_len
|
||||
self.rope_theta = rope_theta
|
||||
self.dropout = dropout
|
||||
self.flash_attn = flash_attn
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
self.use_moe = use_moe
|
||||
self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
|
||||
self.n_routed_experts = n_routed_experts # 总的专家数量
|
||||
self.n_shared_experts = n_shared_experts # 共享专家
|
||||
self.scoring_func = scoring_func # 评分函数，默认为'softmax'
|
||||
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
|
||||
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
|
||||
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
|
||||
super().__init__(**kwargs)
|
||||
|
|
@ -1,16 +0,0 @@
|
|||
from .LMConfig import LMConfig
|
||||
from typing import List
|
||||
|
||||
|
||||
class VLMConfig(LMConfig):
|
||||
model_type = "minimind-v"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image_special_token: str = '@' * 196,
|
||||
image_ids: List = [34] * 196,
|
||||
**kwargs,
|
||||
):
|
||||
self.image_special_token = image_special_token
|
||||
self.image_ids = image_ids
|
||||
super().__init__(**kwargs)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
385
model/model.py
385
model/model.py
|
|
@ -1,385 +0,0 @@
|
|||
import math
|
||||
import struct
|
||||
import inspect
|
||||
import time
|
||||
|
||||
from .LMConfig import LMConfig
|
||||
from typing import Any, Optional, Tuple, List, Union
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from transformers import PreTrainedModel
|
||||
from transformers.modeling_outputs import CausalLMOutputWithPast
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self.weight * self._norm(x.float()).type_as(x)
|
||||
|
||||
|
||||
def precompute_pos_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
t = torch.arange(end, device=freqs.device) # type: ignore
|
||||
freqs = torch.outer(t, freqs).float() # type: ignore
|
||||
pos_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
|
||||
return pos_cis
|
||||
|
||||
|
||||
def apply_rotary_emb(xq, xk, pos_cis):
|
||||
def unite_shape(pos_cis, x):
|
||||
ndim = x.ndim
|
||||
assert 0 <= 1 < ndim
|
||||
assert pos_cis.shape == (x.shape[1], x.shape[-1])
|
||||
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
|
||||
return pos_cis.view(*shape)
|
||||
|
||||
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
|
||||
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
|
||||
pos_cis = unite_shape(pos_cis, xq_)
|
||||
xq_out = torch.view_as_real(xq_ * pos_cis).flatten(3)
|
||||
xk_out = torch.view_as_real(xk_ * pos_cis).flatten(3)
|
||||
return xq_out.type_as(xq), xk_out.type_as(xk)
|
||||
|
||||
|
||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
||||
bs, slen, n_kv_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, n_kv_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: LMConfig):
|
||||
super().__init__()
|
||||
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
|
||||
assert args.n_heads % self.n_kv_heads == 0
|
||||
self.n_local_heads = args.n_heads
|
||||
self.n_local_kv_heads = self.n_kv_heads
|
||||
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
||||
self.head_dim = args.dim // args.n_heads
|
||||
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
|
||||
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
|
||||
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
|
||||
self.attn_dropout = nn.Dropout(args.dropout)
|
||||
self.resid_dropout = nn.Dropout(args.dropout)
|
||||
self.dropout = args.dropout
|
||||
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
|
||||
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
||||
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
|
||||
mask = torch.triu(mask, diagonal=1)
|
||||
self.register_buffer("mask", mask, persistent=False)
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
pos_cis: torch.Tensor,
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
use_cache=False):
|
||||
bsz, seq_len, _ = x.shape
|
||||
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
|
||||
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
||||
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
|
||||
xq, xk = apply_rotary_emb(xq, xk, pos_cis)
|
||||
# kv_cache实现
|
||||
if past_key_value is not None:
|
||||
xk = torch.cat([past_key_value[0], xk], dim=1)
|
||||
xv = torch.cat([past_key_value[1], xv], dim=1)
|
||||
past_kv = (xk, xv) if use_cache else None
|
||||
|
||||
xq, xk, xv = (
|
||||
xq.transpose(1, 2),
|
||||
repeat_kv(xk, self.n_rep).transpose(1, 2),
|
||||
repeat_kv(xv, self.n_rep).transpose(1, 2)
|
||||
)
|
||||
if self.flash and seq_len != 1:
|
||||
dropout_p = self.dropout if self.training else 0.0
|
||||
output = F.scaled_dot_product_attention(
|
||||
xq, xk, xv,
|
||||
attn_mask=None,
|
||||
dropout_p=dropout_p,
|
||||
is_causal=True
|
||||
)
|
||||
else:
|
||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||
scores += self.mask[:, :, :seq_len, :seq_len]
|
||||
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
||||
scores = self.attn_dropout(scores)
|
||||
output = scores @ xv
|
||||
|
||||
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
||||
output = self.resid_dropout(self.wo(output))
|
||||
return output, past_kv
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
if config.hidden_dim is None:
|
||||
hidden_dim = 4 * config.dim
|
||||
hidden_dim = int(2 * hidden_dim / 3)
|
||||
config.hidden_dim = config.multiple_of * ((hidden_dim + config.multiple_of - 1) // config.multiple_of)
|
||||
self.w1 = nn.Linear(config.dim, config.hidden_dim, bias=False)
|
||||
self.w2 = nn.Linear(config.hidden_dim, config.dim, bias=False)
|
||||
self.w3 = nn.Linear(config.dim, config.hidden_dim, bias=False)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
|
||||
def forward(self, x):
|
||||
return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
|
||||
self.scoring_func = config.scoring_func
|
||||
self.alpha = config.aux_loss_alpha
|
||||
self.seq_aux = config.seq_aux
|
||||
|
||||
self.norm_topk_prob = config.norm_topk_prob
|
||||
self.gating_dim = config.dim
|
||||
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
import torch.nn.init as init
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
|
||||
def forward(self, hidden_states):
|
||||
bsz, seq_len, h = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, h)
|
||||
logits = F.linear(hidden_states, self.weight, None)
|
||||
if self.scoring_func == 'softmax':
|
||||
scores = logits.softmax(dim=-1)
|
||||
else:
|
||||
raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
|
||||
|
||||
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
|
||||
|
||||
if self.top_k > 1 and self.norm_topk_prob:
|
||||
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
|
||||
topk_weight = topk_weight / denominator
|
||||
|
||||
if self.training and self.alpha > 0.0:
|
||||
scores_for_aux = scores
|
||||
aux_topk = self.top_k
|
||||
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
|
||||
if self.seq_aux:
|
||||
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
|
||||
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
|
||||
ce.scatter_add_(1, topk_idx_for_aux_loss,
|
||||
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
|
||||
seq_len * aux_topk / self.n_routed_experts)
|
||||
aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
|
||||
else:
|
||||
mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
|
||||
ce = mask_ce.float().mean(0)
|
||||
Pi = scores_for_aux.mean(0)
|
||||
fi = ce * self.n_routed_experts
|
||||
aux_loss = (Pi * fi).sum() * self.alpha
|
||||
else:
|
||||
aux_loss = 0
|
||||
return topk_idx, topk_weight, aux_loss
|
||||
|
||||
|
||||
class MOEFeedForward(nn.Module):
|
||||
def __init__(self, config: LMConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_routed_experts)
|
||||
])
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts is not None:
|
||||
self.shared_experts = FeedForward(config)
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
orig_shape = x.shape
|
||||
bsz, seq_len, _ = x.shape
|
||||
# 使用门控机制选择专家
|
||||
topk_idx, topk_weight, aux_loss = self.gate(x)
|
||||
x = x.view(-1, x.shape[-1])
|
||||
flat_topk_idx = topk_idx.view(-1)
|
||||
if self.training:
|
||||
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
|
||||
y = torch.empty_like(x, dtype=torch.float16)
|
||||
for i, expert in enumerate(self.experts):
|
||||
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
|
||||
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
|
||||
y = y.view(*orig_shape)
|
||||
else:
|
||||
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
|
||||
if self.config.n_shared_experts is not None:
|
||||
y = y + self.shared_experts(identity)
|
||||
self.aux_loss = aux_loss
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||
expert_cache = torch.zeros_like(x)
|
||||
idxs = flat_expert_indices.argsort()
|
||||
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
|
||||
token_idxs = idxs // self.config.num_experts_per_tok
|
||||
# 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
|
||||
# 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
|
||||
# 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
|
||||
# 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
|
||||
for i, end_idx in enumerate(tokens_per_expert):
|
||||
start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
|
||||
if start_idx == end_idx:
|
||||
continue
|
||||
expert = self.experts[i]
|
||||
exp_token_idx = token_idxs[start_idx:end_idx]
|
||||
expert_tokens = x[exp_token_idx]
|
||||
expert_out = expert(expert_tokens).to(expert_cache.dtype)
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
|
||||
|
||||
return expert_cache
|
||||
|
||||
|
||||
class MiniMindBlock(nn.Module):
|
||||
def __init__(self, layer_id: int, config: LMConfig):
|
||||
super().__init__()
|
||||
self.n_heads = config.n_heads
|
||||
self.dim = config.dim
|
||||
self.head_dim = config.dim // config.n_heads
|
||||
self.attention = Attention(config)
|
||||
|
||||
self.layer_id = layer_id
|
||||
self.attention_norm = RMSNorm(config.dim, eps=config.norm_eps)
|
||||
self.ffn_norm = RMSNorm(config.dim, eps=config.norm_eps)
|
||||
self.feed_forward = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
||||
|
||||
def forward(self, x, pos_cis, past_key_value=None, use_cache=False):
|
||||
h_attn, past_kv = self.attention(
|
||||
self.attention_norm(x),
|
||||
pos_cis,
|
||||
past_key_value=past_key_value,
|
||||
use_cache=use_cache
|
||||
)
|
||||
h = x + h_attn
|
||||
out = h + self.feed_forward(self.ffn_norm(h))
|
||||
return out, past_kv
|
||||
|
||||
|
||||
class MiniMindLM(PreTrainedModel):
|
||||
config_class = LMConfig
|
||||
|
||||
def __init__(self, params: LMConfig = None):
|
||||
self.params = params or LMConfig()
|
||||
super().__init__(self.params)
|
||||
self.vocab_size, self.n_layers = params.vocab_size, params.n_layers
|
||||
self.tok_embeddings = nn.Embedding(params.vocab_size, params.dim)
|
||||
self.dropout = nn.Dropout(params.dropout)
|
||||
self.layers = nn.ModuleList([MiniMindBlock(l, params) for l in range(self.n_layers)])
|
||||
self.norm = RMSNorm(params.dim, eps=params.norm_eps)
|
||||
self.output = nn.Linear(params.dim, params.vocab_size, bias=False)
|
||||
self.tok_embeddings.weight = self.output.weight
|
||||
self.register_buffer("pos_cis",
|
||||
precompute_pos_cis(dim=params.dim // params.n_heads, theta=params.rope_theta),
|
||||
persistent=False)
|
||||
self.OUT = CausalLMOutputWithPast()
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**args):
|
||||
past_key_values = past_key_values or [None] * len(self.layers)
|
||||
start_pos = args.get('start_pos', 0)
|
||||
h = self.dropout(self.tok_embeddings(input_ids))
|
||||
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.size(1)]
|
||||
past_kvs = []
|
||||
for l, layer in enumerate(self.layers):
|
||||
h, past_kv = layer(
|
||||
h, pos_cis,
|
||||
past_key_value=past_key_values[l],
|
||||
use_cache=use_cache
|
||||
)
|
||||
past_kvs.append(past_kv)
|
||||
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.output(self.norm(h)[:, slice_indices, :])
|
||||
aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
|
||||
self.OUT.__setitem__('last_hidden_state', h)
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
return self.OUT
|
||||
|
||||
@torch.inference_mode()
|
||||
def generate(self, input_ids, eos_token_id=2, max_new_tokens=1024, temperature=0.75, top_p=0.90,
|
||||
stream=False, rp=1., use_cache=True, pad_token_id=0, num_return_sequences=1, **args):
|
||||
# 流式生成
|
||||
if stream:
|
||||
return self._stream(input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
|
||||
|
||||
# 直接生成
|
||||
generated = []
|
||||
for i in range(input_ids.size(0)):
|
||||
non_pad = input_ids[i][input_ids[i] != pad_token_id].unsqueeze(0)
|
||||
for _ in range(num_return_sequences):
|
||||
out = self._stream(non_pad, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args)
|
||||
tokens_list = [tokens[:, -1:] for tokens in out]
|
||||
gen = torch.cat(tokens_list, dim=-1) if tokens_list else non_pad
|
||||
full_sequence = torch.cat([non_pad, gen], dim=-1)
|
||||
generated.append(full_sequence)
|
||||
|
||||
max_length = max(seq.size(1) for seq in generated)
|
||||
generated = [
|
||||
torch.cat(
|
||||
[seq, torch.full((1, max_length - seq.size(1)), pad_token_id, dtype=seq.dtype, device=seq.device)],
|
||||
dim=-1)
|
||||
for seq in generated
|
||||
]
|
||||
output = torch.cat(generated, dim=0)
|
||||
res = output.view(input_ids.size(0) * num_return_sequences, -1)
|
||||
return res
|
||||
|
||||
def _stream(self, input_ids, eos_token_id, max_new_tokens, temperature, top_p, rp, use_cache, **args):
|
||||
start, first_seq, past_kvs = input_ids.shape[1], True, None
|
||||
while input_ids.shape[1] < max_new_tokens - 1:
|
||||
if first_seq or not use_cache:
|
||||
out, first_seq = self(input_ids, past_key_values=past_kvs, use_cache=use_cache, **args), False
|
||||
else:
|
||||
out = self(input_ids[:, -1:], past_key_values=past_kvs, use_cache=use_cache,
|
||||
start_pos=input_ids.shape[1] - 1, **args)
|
||||
logits, past_kvs = out.logits[:, -1, :], out.past_key_values
|
||||
logits[:, list(set(input_ids.tolist()[0]))] /= rp
|
||||
logits /= (temperature + 1e-9)
|
||||
if top_p is not None and top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
||||
sorted_probs = F.softmax(sorted_logits, dim=-1)
|
||||
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 = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
|
||||
logits[indices_to_remove] = -float('Inf')
|
||||
input_ids_next = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
|
||||
input_ids = torch.cat((input_ids, input_ids_next), dim=1)
|
||||
yield input_ids[:, start:]
|
||||
if input_ids_next.item() == eos_token_id:
|
||||
break
|
||||
|
|
@ -0,0 +1,446 @@
|
|||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
# MiniMind Config
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
|
||||
class MiniMindConfig(PretrainedConfig):
|
||||
model_type = "minimind"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dropout: float = 0.0,
|
||||
bos_token_id: int = 1,
|
||||
eos_token_id: int = 2,
|
||||
hidden_act: str = 'silu',
|
||||
hidden_size: int = 512,
|
||||
intermediate_size: int = None,
|
||||
max_position_embeddings: int = 32768,
|
||||
num_attention_heads: int = 8,
|
||||
num_hidden_layers: int = 8,
|
||||
num_key_value_heads: int = 2,
|
||||
vocab_size: int = 6400,
|
||||
rms_norm_eps: float = 1e-05,
|
||||
rope_theta: int = 1000000.0,
|
||||
flash_attn: bool = True,
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
use_moe: bool = False,
|
||||
num_experts_per_tok: int = 2,
|
||||
n_routed_experts: int = 4,
|
||||
n_shared_experts: int = 1,
|
||||
scoring_func: str = 'softmax',
|
||||
aux_loss_alpha: float = 0.1,
|
||||
seq_aux: bool = True,
|
||||
norm_topk_prob: bool = True,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.dropout = dropout
|
||||
self.bos_token_id = bos_token_id
|
||||
self.eos_token_id = eos_token_id
|
||||
self.hidden_act = hidden_act
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.max_position_embeddings = max_position_embeddings
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.num_hidden_layers = num_hidden_layers
|
||||
self.num_key_value_heads = num_key_value_heads
|
||||
self.vocab_size = vocab_size
|
||||
self.rms_norm_eps = rms_norm_eps
|
||||
self.rope_theta = rope_theta
|
||||
self.flash_attn = flash_attn
|
||||
####################################################
|
||||
# Here are the specific configurations of MOE
|
||||
# When use_moe is false, the following is invalid
|
||||
####################################################
|
||||
self.use_moe = use_moe
|
||||
self.num_experts_per_tok = num_experts_per_tok # 每个token选择的专家数量
|
||||
self.n_routed_experts = n_routed_experts # 总的专家数量
|
||||
self.n_shared_experts = n_shared_experts # 共享专家
|
||||
self.scoring_func = scoring_func # 评分函数，默认为'softmax'
|
||||
self.aux_loss_alpha = aux_loss_alpha # 辅助损失的alpha参数
|
||||
self.seq_aux = seq_aux # 是否在序列级别上计算辅助损失
|
||||
self.norm_topk_prob = norm_topk_prob # 是否标准化top-k概率
|
||||
|
||||
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
# MiniMind Model
|
||||
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
|
||||
|
||||
import math
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.activations import ACT2FN
|
||||
from typing import Optional, Tuple, List, Union
|
||||
import torch.nn.functional as F
|
||||
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
|
||||
from transformers.modeling_outputs import CausalLMOutputWithPast
|
||||
|
||||
|
||||
class RMSNorm(torch.nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
|
||||
def forward(self, x):
|
||||
return self.weight * self._norm(x.float()).type_as(x)
|
||||
|
||||
|
||||
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
|
||||
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
|
||||
t = torch.arange(end, device=freqs.device)
|
||||
freqs = torch.outer(t, freqs).float()
|
||||
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
|
||||
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
|
||||
return freqs_cos, freqs_sin
|
||||
|
||||
|
||||
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
|
||||
def rotate_half(x):
|
||||
return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
|
||||
|
||||
q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
|
||||
k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
|
||||
return q_embed, k_embed
|
||||
|
||||
|
||||
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
|
||||
bs, slen, num_key_value_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, num_key_value_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, args: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
|
||||
assert args.num_attention_heads % self.num_key_value_heads == 0
|
||||
self.n_local_heads = args.num_attention_heads
|
||||
self.n_local_kv_heads = self.num_key_value_heads
|
||||
self.n_rep = self.n_local_heads // self.n_local_kv_heads
|
||||
self.head_dim = args.hidden_size // args.num_attention_heads
|
||||
self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
|
||||
self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
||||
self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
|
||||
self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False)
|
||||
self.attn_dropout = nn.Dropout(args.dropout)
|
||||
self.resid_dropout = nn.Dropout(args.dropout)
|
||||
self.dropout = args.dropout
|
||||
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
|
||||
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
|
||||
|
||||
def forward(self,
|
||||
x: torch.Tensor,
|
||||
position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 修改为接收cos和sin
|
||||
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
|
||||
use_cache=False,
|
||||
attention_mask: Optional[torch.Tensor] = None):
|
||||
bsz, seq_len, _ = x.shape
|
||||
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
|
||||
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
|
||||
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
|
||||
|
||||
cos, sin = position_embeddings
|
||||
xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
|
||||
|
||||
# kv_cache实现
|
||||
if past_key_value is not None:
|
||||
xk = torch.cat([past_key_value[0], xk], dim=1)
|
||||
xv = torch.cat([past_key_value[1], xv], dim=1)
|
||||
past_kv = (xk, xv) if use_cache else None
|
||||
|
||||
xq, xk, xv = (
|
||||
xq.transpose(1, 2),
|
||||
repeat_kv(xk, self.n_rep).transpose(1, 2),
|
||||
repeat_kv(xv, self.n_rep).transpose(1, 2)
|
||||
)
|
||||
|
||||
if False and self.flash and seq_len != 1:
|
||||
dropout_p = self.dropout if self.training else 0.0
|
||||
attn_mask = None
|
||||
if attention_mask is not None:
|
||||
attn_mask = attention_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_local_heads, seq_len, -1)
|
||||
attn_mask = attn_mask.bool() if attention_mask is not None else None
|
||||
|
||||
output = F.scaled_dot_product_attention(xq, xk, xv, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=True)
|
||||
else:
|
||||
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
|
||||
scores = scores + torch.triu(
|
||||
torch.full((seq_len, seq_len), float("-inf"), device=scores.device),
|
||||
diagonal=1
|
||||
).unsqueeze(0).unsqueeze(0) # scores+mask
|
||||
|
||||
if attention_mask is not None:
|
||||
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
|
||||
extended_attention_mask = (1.0 - extended_attention_mask) * -1e9
|
||||
scores = scores + extended_attention_mask
|
||||
|
||||
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
|
||||
scores = self.attn_dropout(scores)
|
||||
output = scores @ xv
|
||||
|
||||
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
|
||||
output = self.resid_dropout(self.o_proj(output))
|
||||
return output, past_kv
|
||||
|
||||
|
||||
class FeedForward(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
if config.intermediate_size is None:
|
||||
intermediate_size = int(config.hidden_size * 8 / 3)
|
||||
config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64)
|
||||
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
|
||||
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x):
|
||||
return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
|
||||
|
||||
|
||||
class MoEGate(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.top_k = config.num_experts_per_tok
|
||||
self.n_routed_experts = config.n_routed_experts
|
||||
|
||||
self.scoring_func = config.scoring_func
|
||||
self.alpha = config.aux_loss_alpha
|
||||
self.seq_aux = config.seq_aux
|
||||
|
||||
self.norm_topk_prob = config.norm_topk_prob
|
||||
self.gating_dim = config.hidden_size
|
||||
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
import torch.nn.init as init
|
||||
init.kaiming_uniform_(self.weight, a=math.sqrt(5))
|
||||
|
||||
def forward(self, hidden_states):
|
||||
bsz, seq_len, h = hidden_states.shape
|
||||
hidden_states = hidden_states.view(-1, h)
|
||||
logits = F.linear(hidden_states, self.weight, None)
|
||||
if self.scoring_func == 'softmax':
|
||||
scores = logits.softmax(dim=-1)
|
||||
else:
|
||||
raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
|
||||
|
||||
topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
|
||||
|
||||
if self.top_k > 1 and self.norm_topk_prob:
|
||||
denominator = topk_weight.sum(dim=-1, keepdim=True) + 1e-20
|
||||
topk_weight = topk_weight / denominator
|
||||
|
||||
if self.training and self.alpha > 0.0:
|
||||
scores_for_aux = scores
|
||||
aux_topk = self.top_k
|
||||
topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
|
||||
if self.seq_aux:
|
||||
scores_for_seq_aux = scores_for_aux.view(bsz, seq_len, -1)
|
||||
ce = torch.zeros(bsz, self.n_routed_experts, device=hidden_states.device)
|
||||
ce.scatter_add_(1, topk_idx_for_aux_loss,
|
||||
torch.ones(bsz, seq_len * aux_topk, device=hidden_states.device)).div_(
|
||||
seq_len * aux_topk / self.n_routed_experts)
|
||||
aux_loss = (ce * scores_for_seq_aux.mean(dim=1)).sum(dim=1).mean() * self.alpha
|
||||
else:
|
||||
mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
|
||||
ce = mask_ce.float().mean(0)
|
||||
Pi = scores_for_aux.mean(0)
|
||||
fi = ce * self.n_routed_experts
|
||||
aux_loss = (Pi * fi).sum() * self.alpha
|
||||
else:
|
||||
aux_loss = 0
|
||||
return topk_idx, topk_weight, aux_loss
|
||||
|
||||
|
||||
class MOEFeedForward(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_routed_experts)
|
||||
])
|
||||
self.gate = MoEGate(config)
|
||||
if config.n_shared_experts > 0:
|
||||
self.shared_experts = nn.ModuleList([
|
||||
FeedForward(config)
|
||||
for _ in range(config.n_shared_experts)
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
orig_shape = x.shape
|
||||
bsz, seq_len, _ = x.shape
|
||||
# 使用门控机制选择专家
|
||||
topk_idx, topk_weight, aux_loss = self.gate(x)
|
||||
x = x.view(-1, x.shape[-1])
|
||||
flat_topk_idx = topk_idx.view(-1)
|
||||
if self.training:
|
||||
x = x.repeat_interleave(self.config.num_experts_per_tok, dim=0)
|
||||
y = torch.empty_like(x, dtype=torch.float16)
|
||||
for i, expert in enumerate(self.experts):
|
||||
y[flat_topk_idx == i] = expert(x[flat_topk_idx == i]).to(y.dtype) # 确保类型一致
|
||||
y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
|
||||
y = y.view(*orig_shape)
|
||||
else:
|
||||
y = self.moe_infer(x, flat_topk_idx, topk_weight.view(-1, 1)).view(*orig_shape)
|
||||
if self.config.n_shared_experts > 0:
|
||||
for expert in self.shared_experts:
|
||||
y = y + expert(identity)
|
||||
self.aux_loss = aux_loss
|
||||
return y
|
||||
|
||||
@torch.no_grad()
|
||||
def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
|
||||
expert_cache = torch.zeros_like(x)
|
||||
idxs = flat_expert_indices.argsort()
|
||||
tokens_per_expert = flat_expert_indices.bincount().cpu().numpy().cumsum(0)
|
||||
token_idxs = idxs // self.config.num_experts_per_tok
|
||||
# 当tokens_per_expert = [6, 15, 20, 26]，tokens_per_expert.shape[0]即为专家数量（此时为4）
|
||||
# 且token_idxs = [3, 7, 19, 21, 24, 25, 4, 5, 6, 10, 11, 12...] 时
|
||||
# 意味token_idxs[:6] -> [3, 7, 19, 21, 24, 25]这6个位置属于专家0处理的token（每个token有可能被多个专家处理，这取决于num_experts_per_tok）
|
||||
# 接下来9个位置token_idxs[6:15] -> [4, 5, 6, 10, 11, 12...]属于专家1处理的token...依此类推
|
||||
for i, end_idx in enumerate(tokens_per_expert):
|
||||
start_idx = 0 if i == 0 else tokens_per_expert[i - 1]
|
||||
if start_idx == end_idx:
|
||||
continue
|
||||
expert = self.experts[i]
|
||||
exp_token_idx = token_idxs[start_idx:end_idx]
|
||||
expert_tokens = x[exp_token_idx]
|
||||
expert_out = expert(expert_tokens).to(expert_cache.dtype)
|
||||
expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
|
||||
expert_cache.scatter_add_(0, exp_token_idx.view(-1, 1).repeat(1, x.shape[-1]), expert_out)
|
||||
|
||||
return expert_cache
|
||||
|
||||
|
||||
class MiniMindBlock(nn.Module):
|
||||
def __init__(self, layer_id: int, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.hidden_size = config.hidden_size
|
||||
self.head_dim = config.hidden_size // config.num_attention_heads
|
||||
self.self_attn = Attention(config)
|
||||
|
||||
self.layer_id = layer_id
|
||||
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)
|
||||
|
||||
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
|
||||
residual = hidden_states
|
||||
hidden_states, present_key_value = self.self_attn(
|
||||
self.input_layernorm(hidden_states), position_embeddings,
|
||||
past_key_value, use_cache, attention_mask
|
||||
)
|
||||
hidden_states += residual
|
||||
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
|
||||
return hidden_states, present_key_value
|
||||
|
||||
|
||||
class MiniMindModel(nn.Module):
|
||||
def __init__(self, config: MiniMindConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
|
||||
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
|
||||
self.dropout = nn.Dropout(config.dropout)
|
||||
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
|
||||
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
|
||||
end=config.max_position_embeddings, theta=config.rope_theta)
|
||||
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
|
||||
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
**kwargs):
|
||||
batch_size, seq_length = input_ids.shape
|
||||
past_key_values = past_key_values or [None] * len(self.layers)
|
||||
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
||||
|
||||
hidden_states = self.dropout(self.embed_tokens(input_ids))
|
||||
|
||||
position_embeddings = (
|
||||
self.freqs_cos[start_pos:start_pos + seq_length],
|
||||
self.freqs_sin[start_pos:start_pos + seq_length]
|
||||
)
|
||||
|
||||
presents = []
|
||||
for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_key_values)):
|
||||
hidden_states, present = layer(
|
||||
hidden_states,
|
||||
position_embeddings,
|
||||
past_key_value=past_key_value,
|
||||
use_cache=use_cache,
|
||||
attention_mask=attention_mask
|
||||
)
|
||||
presents.append(present)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
aux_loss = sum(
|
||||
layer.mlp.aux_loss
|
||||
for layer in self.layers
|
||||
if isinstance(layer.mlp, MOEFeedForward)
|
||||
)
|
||||
|
||||
return hidden_states, presents, aux_loss
|
||||
|
||||
|
||||
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
|
||||
config_class = MiniMindConfig
|
||||
|
||||
def __init__(self, config: MiniMindConfig = None):
|
||||
self.config = config or MiniMindConfig()
|
||||
super().__init__(self.config)
|
||||
self.model = MiniMindModel(self.config)
|
||||
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
|
||||
self.model.embed_tokens.weight = self.lm_head.weight
|
||||
self.OUT = CausalLMOutputWithPast()
|
||||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**args):
|
||||
h, past_kvs, aux_loss = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
**args
|
||||
)
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.lm_head(h[:, slice_indices, :])
|
||||
self.OUT.__setitem__('last_hidden_state', h)
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
return self.OUT
|
||||
|
|
@ -1,21 +1,36 @@
|
|||
from .VLMConfig import VLMConfig
|
||||
from .model import *
|
||||
import os
|
||||
|
||||
import torch
|
||||
import warnings
|
||||
from .model_minimind import *
|
||||
from typing import Optional, Tuple, List
|
||||
from torch import nn
|
||||
import warnings
|
||||
from transformers import CLIPProcessor, CLIPModel
|
||||
import torch
|
||||
from typing import List
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
class VLMConfig(MiniMindConfig):
|
||||
model_type = "minimind-v"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
image_special_token: str = '@' * 196,
|
||||
image_ids: List = [34] * 196,
|
||||
**kwargs,
|
||||
):
|
||||
self.image_special_token = image_special_token
|
||||
self.image_ids = image_ids
|
||||
super().__init__(**kwargs)
|
||||
|
||||
class VisionProj(nn.Module):
|
||||
def __init__(self, ve_dim=768, lm_dim=512):
|
||||
def __init__(self, ve_hidden_size=768, hidden_size=512):
|
||||
super().__init__()
|
||||
self.ve_dim = ve_dim
|
||||
self.lm_dim = lm_dim
|
||||
self.ve_hidden_size = ve_hidden_size
|
||||
self.hidden_size = hidden_size
|
||||
self.vision_proj = nn.Sequential(
|
||||
nn.Linear(self.ve_dim, self.lm_dim)
|
||||
nn.Linear(self.ve_hidden_size, self.hidden_size)
|
||||
)
|
||||
|
||||
def forward(self, image_encoders):
|
||||
|
|
@ -24,18 +39,22 @@ class VisionProj(nn.Module):
|
|||
|
||||
|
||||
# 继承自语言模型
|
||||
class MiniMindVLM(MiniMindLM):
|
||||
class MiniMindVLM(MiniMindForCausalLM):
|
||||
config_class = VLMConfig
|
||||
|
||||
def __init__(self, params: VLMConfig = None):
|
||||
def __init__(self, params: VLMConfig = None, vision_model_path="./model/vision_model/clip-vit-base-patch16"):
|
||||
super().__init__(params)
|
||||
if not params: params = VLMConfig()
|
||||
self.params = params
|
||||
self.vision_encoder, self.processor = self.__class__.get_vision_model()
|
||||
self.vision_proj = VisionProj(lm_dim=params.dim)
|
||||
self.vision_encoder, self.processor = self.__class__.get_vision_model(vision_model_path)
|
||||
self.vision_proj = VisionProj(hidden_size=params.hidden_size)
|
||||
|
||||
@staticmethod
|
||||
def get_vision_model(model_path="./model/vision_model/clip-vit-base-patch16"):
|
||||
def get_vision_model(model_path: str):
|
||||
from transformers import logging as hf_logging
|
||||
hf_logging.set_verbosity_error()
|
||||
if not os.path.exists(model_path):
|
||||
return None, None
|
||||
model = CLIPModel.from_pretrained(model_path)
|
||||
processor = CLIPProcessor.from_pretrained(model_path)
|
||||
# 冻结 vision_encoder 的所有参数
|
||||
|
|
@ -93,38 +112,57 @@ class MiniMindVLM(MiniMindLM):
|
|||
|
||||
def forward(self,
|
||||
input_ids: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
|
||||
use_cache: bool = False,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
pixel_values: Optional[torch.FloatTensor] = None,
|
||||
**args):
|
||||
start_pos = args.get('start_pos', 0)
|
||||
pixel_tensors = args.get('pixel_tensors', None)
|
||||
h = self.tok_embeddings(input_ids)
|
||||
batch_size, seq_length = input_ids.shape
|
||||
past_key_values = past_key_values or [None] * len(self.model.layers)
|
||||
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
|
||||
|
||||
if pixel_tensors is not None and start_pos == 0:
|
||||
if len(pixel_tensors.shape) == 6:
|
||||
pixel_tensors = pixel_tensors.squeeze(2)
|
||||
bs, num, c, im_h, im_w = pixel_tensors.shape
|
||||
hidden_states = self.model.dropout(self.model.embed_tokens(input_ids))
|
||||
|
||||
if pixel_values is not None and start_pos == 0:
|
||||
if len(pixel_values.shape) == 6:
|
||||
pixel_values = pixel_values.squeeze(2)
|
||||
bs, num, c, im_h, im_w = pixel_values.shape
|
||||
stack_dim = 1 if bs > 1 else 0
|
||||
vision_tensors = torch.stack([
|
||||
MiniMindVLM.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.vision_encoder)
|
||||
MiniMindVLM.get_image_embeddings(pixel_values[:, i, :, :, :], self.vision_encoder)
|
||||
for i in range(num)
|
||||
], dim=stack_dim)
|
||||
h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
|
||||
hidden_states = self.count_vision_proj(tokens=input_ids, h=hidden_states, vision_tensors=vision_tensors,
|
||||
seqlen=input_ids.shape[1])
|
||||
|
||||
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.shape[1]]
|
||||
past_kvs = []
|
||||
for l, layer in enumerate(self.layers):
|
||||
h, past_kv = layer(
|
||||
h, pos_cis,
|
||||
past_key_value=past_key_values[l] if past_key_values else None,
|
||||
use_cache=use_cache
|
||||
position_embeddings = (
|
||||
self.model.freqs_cos[start_pos:start_pos + seq_length],
|
||||
self.model.freqs_sin[start_pos:start_pos + seq_length]
|
||||
)
|
||||
|
||||
presents = []
|
||||
for layer_idx, (layer, past_key_value) in enumerate(zip(self.model.layers, past_key_values)):
|
||||
hidden_states, present = layer(
|
||||
hidden_states,
|
||||
position_embeddings,
|
||||
past_key_value=past_key_value,
|
||||
use_cache=use_cache,
|
||||
attention_mask=attention_mask
|
||||
)
|
||||
past_kvs.append(past_kv)
|
||||
presents.append(present)
|
||||
|
||||
logits = self.output(self.norm(h))
|
||||
aux_loss = sum(l.feed_forward.aux_loss for l in self.layers if isinstance(l.feed_forward, MOEFeedForward))
|
||||
hidden_states = self.model.norm(hidden_states)
|
||||
|
||||
aux_loss = sum(
|
||||
layer.mlp.aux_loss
|
||||
for layer in self.model.layers
|
||||
if isinstance(layer.mlp, MOEFeedForward)
|
||||
)
|
||||
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
|
||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||
self.OUT.__setitem__('last_hidden_state', hidden_states)
|
||||
self.OUT.__setitem__('logits', logits)
|
||||
self.OUT.__setitem__('aux_loss', aux_loss)
|
||||
self.OUT.__setitem__('past_key_values', past_kvs)
|
||||
self.OUT.__setitem__('past_key_values', presents)
|
||||
return self.OUT
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load Diff
|
|
@ -4,7 +4,7 @@
|
|||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "<unk>",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -12,7 +12,7 @@
|
|||
"special": true
|
||||
},
|
||||
"1": {
|
||||
"content": "<s>",
|
||||
"content": "<|im_start|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -20,7 +20,7 @@
|
|||
"special": true
|
||||
},
|
||||
"2": {
|
||||
"content": "</s>",
|
||||
"content": "<|im_end|>",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
|
|
@ -29,15 +29,15 @@
|
|||
}
|
||||
},
|
||||
"additional_special_tokens": [],
|
||||
"bos_token": "<s>",
|
||||
"bos_token": "<|im_start|>",
|
||||
"clean_up_tokenization_spaces": false,
|
||||
"eos_token": "</s>",
|
||||
"eos_token": "<|im_end|>",
|
||||
"legacy": true,
|
||||
"model_max_length": 32768,
|
||||
"pad_token": "<unk>",
|
||||
"pad_token": "<|endoftext|>",
|
||||
"sp_model_kwargs": {},
|
||||
"spaces_between_special_tokens": false,
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "<unk>",
|
||||
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<s>system\\n' + system_message + '</s>\\n' }}{% else %}{{ '<s>system\\n你是 MiniMind，是一个有用的人工智能助手。</s>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<s>user\\n' + content + '</s>\\n<s>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '</s>' + '\\n' }}{% endif %}{% endfor %}"
|
||||
"unk_token": "<|endoftext|>",
|
||||
"chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{{ '<|im_start|>system\\n' + system_message + '<|im_end|>\\n' }}{% else %}{{ '<|im_start|>system\\nYou are a helpful assistant<|im_end|>\\n' }}{% endif %}{% for message in messages %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|im_start|>user\\n' + content + '<|im_end|>\\n<|im_start|>assistant\\n' }}{% elif message['role'] == 'assistant' %}{{ content + '<|im_end|>' + '\\n' }}{% endif %}{% endfor %}"
|
||||
}
|
||||
|
|
@ -0,0 +1,44 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import torch
|
||||
import warnings
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
|
||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||
|
||||
warnings.filterwarnings('ignore', category=UserWarning)
|
||||
|
||||
|
||||
def convert_torch2transformers_minimind(torch_path, transformers_path, dtype=torch.bfloat16):
|
||||
VLMConfig.register_for_auto_class()
|
||||
MiniMindVLM.register_for_auto_class("AutoModelForCausalLM")
|
||||
lm_model = MiniMindVLM(lm_config, vision_model_path="../model/vision_model/clip-vit-base-patch16")
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
state_dict = torch.load(torch_path, map_location=device)
|
||||
lm_model.load_state_dict(state_dict, strict=False)
|
||||
lm_model = lm_model.to(dtype) # 转换模型权重精度
|
||||
model_params = sum(p.numel() for p in lm_model.parameters() if p.requires_grad)
|
||||
print(f'模型参数: {model_params / 1e6} 百万 = {model_params / 1e9} B (Billion)')
|
||||
del lm_model.vision_encoder
|
||||
lm_model.save_pretrained(transformers_path, safe_serialization=False)
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model/')
|
||||
tokenizer.save_pretrained(transformers_path)
|
||||
print(f"模型已保存为 Transformers-MiniMind-V 格式: {transformers_path}")
|
||||
|
||||
|
||||
def convert_transformers2torch(transformers_path, torch_path):
|
||||
model = AutoModelForCausalLM.from_pretrained(transformers_path, trust_remote_code=True)
|
||||
torch.save(model.state_dict(), torch_path)
|
||||
print(f"模型已保存为 PyTorch 格式: {torch_path}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
lm_config = VLMConfig(hidden_size=768, num_hidden_layers=16, max_seq_len=8192, use_moe=False)
|
||||
|
||||
torch_path = f"../out/sft_vlm_{lm_config.hidden_size}{'_moe' if lm_config.use_moe else ''}.pth"
|
||||
|
||||
transformers_path = '../MiniMind2-V'
|
||||
|
||||
convert_torch2transformers_minimind(torch_path, transformers_path)
|
||||
|
|
@ -1,11 +1,17 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
__package__ = "scripts"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
import argparse
|
||||
import torch
|
||||
import warnings
|
||||
import gradio as gr
|
||||
from queue import Queue
|
||||
from threading import Thread
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model_vlm import MiniMindVLM
|
||||
from model.VLMConfig import VLMConfig
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
|
||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||
from transformers import logging as hf_logging
|
||||
|
||||
hf_logging.set_verbosity_error()
|
||||
|
|
@ -13,29 +19,42 @@ warnings.filterwarnings('ignore')
|
|||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model')
|
||||
|
||||
if args.load == 0:
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/sft_vlm_{lm_config.dim}{moe_path}.pth'
|
||||
model = MiniMindVLM(lm_config)
|
||||
ckp = f'../out/sft_vlm_{lm_config.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=args.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-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.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} 百万')
|
||||
|
||||
vision_model, preprocess = MiniMindVLM.get_vision_model()
|
||||
vision_model, preprocess = model.vision_encoder, model.processor
|
||||
return model.eval().to(args.device), tokenizer, vision_model.to(args.device), preprocess
|
||||
|
||||
|
||||
class CustomStreamer(TextStreamer):
|
||||
def __init__(self, tokenizer, queue):
|
||||
super().__init__(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
self.queue = queue
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
def on_finalized_text(self, text: str, stream_end: bool = False):
|
||||
self.queue.put(text)
|
||||
if stream_end:
|
||||
self.queue.put(None)
|
||||
|
||||
|
||||
def chat(prompt, current_image_path):
|
||||
global temperature, top_p
|
||||
image = Image.open(current_image_path).convert('RGB')
|
||||
pixel_tensors = MiniMindVLM.image2tensor(image, preprocess).to(model.device).unsqueeze(0)
|
||||
pixel_values = MiniMindVLM.image2tensor(image, preprocess).to(model.device).unsqueeze(0)
|
||||
|
||||
prompt = f'{lm_config.image_special_token}\n{prompt}'
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
|
@ -45,32 +64,37 @@ def chat(prompt, current_image_path):
|
|||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-args.max_seq_len + 1:]
|
||||
|
||||
with torch.no_grad():
|
||||
x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=model.device).unsqueeze(0)
|
||||
outputs = model.generate(
|
||||
x,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
max_new_tokens=args.max_seq_len,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
stream=True,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
pixel_tensors=pixel_tensors
|
||||
)
|
||||
try:
|
||||
if not args.stream:
|
||||
print(tokenizer.decode(outputs.squeeze()[x.shape[1]:].tolist(), skip_special_tokens=True), end='')
|
||||
else:
|
||||
history_idx = 0
|
||||
for y in outputs:
|
||||
answer = tokenizer.decode(y[0].tolist(), skip_special_tokens=True)
|
||||
if (answer and answer[-1] == '<EFBFBD>') or not answer:
|
||||
continue
|
||||
yield answer[history_idx:]
|
||||
history_idx = len(answer)
|
||||
except StopIteration:
|
||||
print("No answer")
|
||||
print('\n')
|
||||
inputs = tokenizer(
|
||||
new_prompt,
|
||||
return_tensors="pt",
|
||||
truncation=True
|
||||
).to(args.device)
|
||||
queue = Queue()
|
||||
streamer = CustomStreamer(tokenizer, queue)
|
||||
|
||||
def _generate():
|
||||
model.generate(
|
||||
inputs.input_ids,
|
||||
max_new_tokens=args.max_seq_len,
|
||||
do_sample=True,
|
||||
temperature=temperature,
|
||||
top_p=top_p,
|
||||
attention_mask=inputs.attention_mask,
|
||||
pad_token_id=tokenizer.pad_token_id,
|
||||
eos_token_id=tokenizer.eos_token_id,
|
||||
streamer=streamer,
|
||||
pixel_values=pixel_values
|
||||
)
|
||||
|
||||
Thread(target=_generate).start()
|
||||
|
||||
while True:
|
||||
text = queue.get()
|
||||
if text is None:
|
||||
break
|
||||
yield text
|
||||
|
||||
|
||||
def launch_gradio_server(server_name="0.0.0.0", server_port=7788):
|
||||
|
|
@ -163,14 +187,15 @@ if __name__ == '__main__':
|
|||
parser.add_argument('--temperature', default=0.65, type=float)
|
||||
parser.add_argument('--top_p', default=0.85, type=float)
|
||||
parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
|
||||
parser.add_argument('--dim', default=768, type=int)
|
||||
parser.add_argument('--n_layers', default=16, type=int)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
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=0, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
parser.add_argument('--load', default=1, type=int, help="0: 原生torch权重，1: transformers加载")
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||
lm_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len, use_moe=args.use_moe)
|
||||
model, tokenizer, vision_model, preprocess = init_model(lm_config)
|
||||
launch_gradio_server(server_name="0.0.0.0", server_port=8888)
|
||||
|
|
@ -1,26 +1,24 @@
|
|||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import json
|
||||
|
||||
import pandas as pd
|
||||
warnings.filterwarnings('ignore')
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
from model.model_vlm import MiniMindVLM
|
||||
from model.VLMConfig import VLMConfig
|
||||
from model.dataset import VLMDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||
from dataset.lm_dataset import VLMDataset
|
||||
|
||||
|
||||
def Logger(content):
|
||||
|
|
@ -35,17 +33,17 @@ def get_lr(current_step, total_steps, lr):
|
|||
def train_epoch(epoch, wandb):
|
||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||
start_time = time.time()
|
||||
for step, (X, Y, loss_mask, pixel_tensors) in enumerate(train_loader):
|
||||
for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
pixel_tensors = pixel_tensors.to(args.device)
|
||||
pixel_values = pixel_values.to(args.device)
|
||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
with ctx:
|
||||
res = model(X, pixel_tensors=pixel_tensors)
|
||||
res = model(X, pixel_values=pixel_values)
|
||||
loss = loss_fct(
|
||||
res.logits.view(-1, res.logits.size(-1)),
|
||||
Y.view(-1)
|
||||
|
|
@ -86,7 +84,7 @@ 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 model_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
|
|
@ -94,16 +92,17 @@ def train_epoch(epoch, wandb):
|
|||
clean_state_dict = {
|
||||
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
|
||||
}
|
||||
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
||||
torch.save(clean_state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(model_config: VLMConfig):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model', use_fast=True)
|
||||
moe_path = '_moe' if model_config.use_moe else ''
|
||||
# 加载纯语言模型权重
|
||||
ckp = f'./out/lm_{model_config.dim}{moe_path}.pth'
|
||||
model = MiniMindVLM(model_config)
|
||||
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)
|
||||
|
||||
|
|
@ -111,16 +110,10 @@ def init_model(model_config: VLMConfig):
|
|||
for name, param in model.named_parameters():
|
||||
if 'vision_proj' not in name:
|
||||
param.requires_grad = False
|
||||
# 可训练
|
||||
if hasattr(model, "layers"):
|
||||
last_two_layers = model.layers[-1:]
|
||||
for layer in last_two_layers:
|
||||
for param in layer.parameters():
|
||||
param.requires_grad = True
|
||||
|
||||
Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
|
||||
_, preprocess = MiniMindVLM.get_vision_model()
|
||||
_, preprocess = model.vision_encoder, model.processor
|
||||
return model.to(args.device), tokenizer, preprocess
|
||||
|
||||
|
||||
|
|
@ -138,7 +131,7 @@ def init_distributed_mode():
|
|||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=4)
|
||||
parser.add_argument("--batch_size", type=int, default=16)
|
||||
parser.add_argument("--learning_rate", type=float, default=4e-4)
|
||||
|
|
@ -147,8 +140,8 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--use_wandb", default=False, action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
|
||||
parser.add_argument("--num_workers", type=int, default=8)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/pretrain_data.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="./dataset/pretrain_images")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_data.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="../dataset/pretrain_images")
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
|
|
@ -156,13 +149,14 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--log_interval", type=int, default=100)
|
||||
parser.add_argument("--save_interval", type=int, default=100)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=640, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len)
|
||||
model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len)
|
||||
max_seq_len = model_config.max_seq_len
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
|
|
@ -1,24 +1,22 @@
|
|||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import json
|
||||
|
||||
import pandas as pd
|
||||
import os
|
||||
import sys
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
__package__ = "trainer"
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
from contextlib import nullcontext
|
||||
from torch import optim, nn
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
from model.model_vlm import MiniMindVLM
|
||||
from model.VLMConfig import VLMConfig
|
||||
from model.dataset import *
|
||||
from model.model_vlm import MiniMindVLM, VLMConfig
|
||||
from dataset.lm_dataset import VLMDataset
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -35,17 +33,17 @@ def get_lr(current_step, total_steps, lr):
|
|||
def train_epoch(epoch, wandb):
|
||||
loss_fct = nn.CrossEntropyLoss(reduction='none')
|
||||
start_time = time.time()
|
||||
for step, (X, Y, loss_mask, pixel_tensors) in enumerate(train_loader):
|
||||
for step, (X, Y, loss_mask, pixel_values) in enumerate(train_loader):
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
pixel_tensors = pixel_tensors.to(args.device)
|
||||
pixel_values = pixel_values.to(args.device)
|
||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch, args.learning_rate)
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
with ctx:
|
||||
res = model(X, pixel_tensors=pixel_tensors)
|
||||
res = model(X, pixel_values=pixel_values)
|
||||
loss = loss_fct(
|
||||
res.logits.view(-1, res.logits.size(-1)),
|
||||
Y.view(-1)
|
||||
|
|
@ -86,7 +84,7 @@ 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 model_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/sft_vlm_{model_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/sft_vlm_{model_config.hidden_size}{moe_path}.pth'
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
|
|
@ -94,23 +92,29 @@ def train_epoch(epoch, wandb):
|
|||
clean_state_dict = {
|
||||
key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
|
||||
}
|
||||
clean_state_dict = {k: v.half() for k, v in clean_state_dict.items()} # 半精度保存
|
||||
torch.save(clean_state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(model_config: VLMConfig):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
tokenizer = AutoTokenizer.from_pretrained('../model')
|
||||
moe_path = '_moe' if model_config.use_moe else ''
|
||||
ckp = f'./out/pretrain_vlm_{model_config.dim}{moe_path}.pth'
|
||||
ckp = f'{args.save_dir}/pretrain_vlm_{model_config.hidden_size}{moe_path}.pth'
|
||||
|
||||
model = MiniMindVLM(model_config)
|
||||
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 = model.to(args.device)
|
||||
|
||||
# 默认全参训练
|
||||
# # 只解冻注意力机制中的投影层参数
|
||||
# for name, param in model.model.named_parameters():
|
||||
# if any(proj in name for proj in ['q_proj', 'k_proj', 'v_proj', 'o_proj']):
|
||||
# param.requires_grad = True
|
||||
|
||||
Logger(f'VLM可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
|
||||
|
||||
_, preprocess = MiniMindVLM.get_vision_model()
|
||||
_, preprocess = model.vision_encoder, model.processor
|
||||
return model.to(args.device), tokenizer, preprocess
|
||||
|
||||
|
||||
|
|
@ -128,17 +132,17 @@ def init_distributed_mode():
|
|||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind-V Pretrain")
|
||||
parser.add_argument("--out_dir", type=str, default="out")
|
||||
parser.add_argument("--epochs", type=int, default=6)
|
||||
parser.add_argument("--batch_size", type=int, default=8)
|
||||
parser.add_argument("--out_dir", type=str, default="../out")
|
||||
parser.add_argument("--epochs", type=int, default=2)
|
||||
parser.add_argument("--batch_size", type=int, default=4)
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-6)
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16")
|
||||
parser.add_argument("--use_wandb", default=False, action="store_true")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V")
|
||||
parser.add_argument("--num_workers", type=int, default=8)
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/sft_data.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="./dataset/sft_images")
|
||||
parser.add_argument("--data_path", type=str, default="../dataset/sft_data.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="../dataset/sft_images")
|
||||
parser.add_argument("--ddp", action="store_true")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1)
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0)
|
||||
|
|
@ -146,13 +150,14 @@ if __name__ == "__main__":
|
|||
parser.add_argument("--log_interval", type=int, default=10)
|
||||
parser.add_argument("--save_interval", type=int, default=10)
|
||||
parser.add_argument('--local_rank', type=int, default=-1)
|
||||
parser.add_argument('--dim', default=512, type=int)
|
||||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--hidden_size', default=512, type=int)
|
||||
parser.add_argument('--num_hidden_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=1536, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len)
|
||||
model_config = VLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers,
|
||||
max_seq_len=args.max_seq_len)
|
||||
max_seq_len = model_config.max_seq_len
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
Loading…
Reference in New Issue