add multi-images
xinyanghuang7
parent
af29f48879
commit
d28564c215
|
|
@ -0,0 +1,2 @@
|
|||
model/clip_model/clip-vit-base-patch32/
|
||||
out/
|
||||
|
|
@ -0,0 +1,244 @@
|
|||
import os
|
||||
import platform
|
||||
import argparse
|
||||
import time
|
||||
import math
|
||||
import warnings
|
||||
import json
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from contextlib import nullcontext
|
||||
|
||||
from torch import optim
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from torch.utils.data import DataLoader, DistributedSampler
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
from model.model import Transformer
|
||||
from model.LMConfig import LMConfig
|
||||
from model.dataset import SFTDataset_multi
|
||||
from model.vision_utils import get_vision_model, get_img_embedding
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def count_parameters(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
def Logger(content):
|
||||
if not ddp or dist.get_rank() == 0:
|
||||
print(content)
|
||||
|
||||
|
||||
def get_lr(it, all):
|
||||
warmup_iters = args.warmup_iters
|
||||
lr_decay_iters = all
|
||||
min_lr = args.learning_rate / 10
|
||||
|
||||
if it < warmup_iters:
|
||||
return args.learning_rate * it / warmup_iters
|
||||
if it > lr_decay_iters:
|
||||
return min_lr
|
||||
decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)
|
||||
assert 0 <= decay_ratio <= 1
|
||||
coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
|
||||
return min_lr + coeff * (args.learning_rate - min_lr)
|
||||
|
||||
|
||||
def train_epoch(epoch, wandb):
|
||||
start_time = time.time()
|
||||
for step, (X, Y, loss_mask, image_process) in enumerate(train_loader):
|
||||
X = X.to(args.device)
|
||||
Y = Y.to(args.device)
|
||||
loss_mask = loss_mask.to(args.device)
|
||||
image_process = image_process.to(args.device)
|
||||
image_encoders = get_img_embedding(image_process, vision_model) # torch.Size([32, 4, 50, 768])
|
||||
lr = get_lr(epoch * iter_per_epoch + step, args.epochs * iter_per_epoch)
|
||||
for param_group in optimizer.param_groups:
|
||||
param_group['lr'] = lr
|
||||
|
||||
with ctx:
|
||||
logits = model(X, Y, image_encoders=image_encoders).logits
|
||||
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), Y.view(-1), ignore_index=0, reduction='none')
|
||||
loss_mask = loss_mask.view(-1)
|
||||
loss = torch.sum(loss * loss_mask) / loss_mask.sum()
|
||||
|
||||
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)
|
||||
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
|
||||
if step % args.log_interval == 0:
|
||||
spend_time = time.time() - start_time
|
||||
Logger(
|
||||
'Epoch:[{}/{}]({}/{}) loss:{:.3f} lr:{:.7f} epoch_Time:{}min:'.format(
|
||||
epoch,
|
||||
args.epochs,
|
||||
step,
|
||||
iter_per_epoch,
|
||||
loss.item(),
|
||||
optimizer.param_groups[-1]['lr'],
|
||||
spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60))
|
||||
|
||||
if (wandb is not None) and (not ddp or dist.get_rank() == 0):
|
||||
wandb.log({"loss": loss,
|
||||
"lr": optimizer.param_groups[-1]['lr'],
|
||||
"epoch_Time": spend_time / (step + 1) * iter_per_epoch // 60 - spend_time // 60})
|
||||
|
||||
if (step + 1) % args.save_interval == 0 and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft.pth'
|
||||
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
if args.save_last and (not ddp or dist.get_rank() == 0):
|
||||
model.eval()
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'{args.save_dir}/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
|
||||
|
||||
if isinstance(model, torch.nn.parallel.DistributedDataParallel):
|
||||
state_dict = model.module.state_dict()
|
||||
else:
|
||||
state_dict = model.state_dict()
|
||||
|
||||
torch.save(state_dict, ckp)
|
||||
model.train()
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
# ckp = f'./out/{lm_config.dim}{moe_path}_vlm_pretrain.pth'
|
||||
ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft.pth'
|
||||
|
||||
model = Transformer(lm_config)
|
||||
state_dict = torch.load(ckp, map_location=args.device)
|
||||
|
||||
# 处理不需要的前缀
|
||||
unwanted_prefix = '_orig_mod.'
|
||||
for k, v in list(state_dict.items()):
|
||||
if k.startswith(unwanted_prefix):
|
||||
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
|
||||
|
||||
# 加载到模型中
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
model = model.to(args.device)
|
||||
|
||||
print(f'模型可学习参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
|
||||
|
||||
(vision_model, preprocess) = get_vision_model()
|
||||
vision_model = vision_model.to(args.device)
|
||||
return model, tokenizer, (vision_model, preprocess)
|
||||
|
||||
|
||||
def init_distributed_mode():
|
||||
if not ddp: return
|
||||
global ddp_local_rank, DEVICE
|
||||
|
||||
dist.init_process_group(backend="nccl")
|
||||
ddp_rank = int(os.environ["RANK"])
|
||||
ddp_local_rank = int(os.environ["LOCAL_RANK"])
|
||||
ddp_world_size = int(os.environ["WORLD_SIZE"])
|
||||
DEVICE = f"cuda:{ddp_local_rank}"
|
||||
torch.cuda.set_device(DEVICE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="MiniMind-V-SFT")
|
||||
parser.add_argument("--out_dir", type=str, default="out", help="Output directory")
|
||||
parser.add_argument("--epochs", type=int, default=19, help="Number of epochs")
|
||||
parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
|
||||
parser.add_argument("--learning_rate", type=float, default=1e-6, help="Learning rate")
|
||||
parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu",
|
||||
help="Device to use")
|
||||
parser.add_argument("--dtype", type=str, default="bfloat16", help="Data type")
|
||||
parser.add_argument("--use_wandb", default=False, action="store_true", help="Use Weights & Biases")
|
||||
parser.add_argument("--wandb_project", type=str, default="MiniMind-V", help="Weights & Biases project name")
|
||||
parser.add_argument("--num_workers", type=int, default=8, help="Number of workers for data loading")
|
||||
parser.add_argument("--data_path", type=str, default="./dataset/sft_multi_images/output_cn.json",
|
||||
help="Path to training data")
|
||||
parser.add_argument("--ddp", action="store_true", help="Use DistributedDataParallel")
|
||||
parser.add_argument("--accumulation_steps", type=int, default=1, help="Gradient accumulation steps")
|
||||
parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping threshold")
|
||||
parser.add_argument("--warmup_iters", type=int, default=0, help="Number of warmup iterations")
|
||||
parser.add_argument("--log_interval", type=int, default=100, help="Logging interval")
|
||||
parser.add_argument("--save_interval", type=int, default=100, help="Model saving interval")
|
||||
parser.add_argument('--local_rank', type=int, default=-1, help='local rank for distributed training')
|
||||
parser.add_argument('--save_last', type=bool, default=True, help='save last step model')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
lm_config = LMConfig()
|
||||
max_seq_len = lm_config.max_seq_len
|
||||
args.save_dir = os.path.join(args.out_dir)
|
||||
os.makedirs(args.save_dir, exist_ok=True)
|
||||
os.makedirs(args.out_dir, exist_ok=True)
|
||||
tokens_per_iter = args.batch_size * max_seq_len
|
||||
torch.manual_seed(1337)
|
||||
device_type = "cuda" if "cuda" in args.device else "cpu"
|
||||
|
||||
args.wandb_run_name = f"MiniMind-V-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
|
||||
|
||||
ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast()
|
||||
ddp = int(os.environ.get("RANK", -1)) != -1 # is this a ddp run?
|
||||
ddp_local_rank, DEVICE = 0, "cuda:0"
|
||||
if ddp:
|
||||
init_distributed_mode()
|
||||
args.device = torch.device(DEVICE)
|
||||
|
||||
if args.use_wandb and (not ddp or ddp_local_rank == 0):
|
||||
import wandb
|
||||
|
||||
wandb.init(project=args.wandb_project, name=args.wandb_run_name)
|
||||
else:
|
||||
wandb = None
|
||||
|
||||
model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
|
||||
|
||||
use_version = 0
|
||||
train_ds = SFTDataset_multi(args.data_path, tokenizer, vision_model=(vision_model, preprocess),
|
||||
image_special_token=lm_config.image_special_token,
|
||||
max_length=max_seq_len)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
batch_size=args.batch_size,
|
||||
pin_memory=True,
|
||||
drop_last=False,
|
||||
shuffle=False,
|
||||
num_workers=args.num_workers,
|
||||
sampler=train_sampler
|
||||
)
|
||||
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype in ['float16', 'bfloat16']))
|
||||
optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
|
||||
|
||||
if False and not lm_config.use_moe and platform.system() != 'Windows' and float(
|
||||
torch.__version__.split('.')[0]) >= 2:
|
||||
Logger("compiling the model... (takes a ~minute)")
|
||||
unoptimized_model = model
|
||||
model = torch.compile(model)
|
||||
|
||||
if ddp:
|
||||
model._ddp_params_and_buffers_to_ignore = {"pos_cis"}
|
||||
model = DistributedDataParallel(model, device_ids=[ddp_local_rank])
|
||||
|
||||
iter_per_epoch = len(train_loader)
|
||||
for epoch in range(args.epochs):
|
||||
train_epoch(epoch, wandb)
|
||||
|
|
@ -0,0 +1,170 @@
|
|||
import os
|
||||
import random
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import warnings
|
||||
|
||||
from PIL import Image
|
||||
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||||
from model.model import Transformer
|
||||
from model.LMConfig import LMConfig
|
||||
from model.vision_utils import get_vision_model, get_img_process, get_img_embedding
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
||||
def count_parameters(model):
|
||||
return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
|
||||
|
||||
def init_model(lm_config):
|
||||
tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
|
||||
model_from = 1 # 1从权重，2用transformers
|
||||
|
||||
if model_from == 1:
|
||||
moe_path = '_moe' if lm_config.use_moe else ''
|
||||
ckp = f'./out/{lm_config.dim}{moe_path}_vlm_sft_multi.pth'
|
||||
|
||||
model = Transformer(lm_config)
|
||||
state_dict = torch.load(ckp, map_location=device)
|
||||
|
||||
# 处理不需要的前缀
|
||||
unwanted_prefix = '_orig_mod.'
|
||||
for k, v in list(state_dict.items()):
|
||||
if k.startswith(unwanted_prefix):
|
||||
state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
|
||||
|
||||
# 加载到模型中
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
else:
|
||||
model = AutoModelForCausalLM.from_pretrained('minimind-v-v1-small',
|
||||
trust_remote_code=True)
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
print(f'模型参数: {count_parameters(model) / 1e6} 百万 = {count_parameters(model) / 1e9} B (Billion)')
|
||||
|
||||
(vision_model, preprocess) = get_vision_model()
|
||||
vision_model = vision_model.to(device)
|
||||
return model, tokenizer, (vision_model, preprocess)
|
||||
|
||||
|
||||
def setup_seed(seed):
|
||||
random.seed(seed) # 设置 Python 的随机种子
|
||||
np.random.seed(seed) # 设置 NumPy 的随机种子
|
||||
torch.manual_seed(seed) # 设置 PyTorch 的随机种子
|
||||
torch.cuda.manual_seed(seed) # 为当前 GPU 设置随机种子（如果有）
|
||||
torch.cuda.manual_seed_all(seed) # 为所有 GPU 设置随机种子（如果有）
|
||||
torch.backends.cudnn.deterministic = True # 确保每次返回的卷积算法是确定的
|
||||
torch.backends.cudnn.benchmark = False # 关闭 cuDNN 的自动调优，避免不确定性
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# -----------------------------------------------------------------------------
|
||||
out_dir = 'out'
|
||||
start = ""
|
||||
temperature = 0.5
|
||||
top_k = 8
|
||||
setup_seed(1337)
|
||||
# device = 'cpu'
|
||||
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
|
||||
dtype = 'bfloat16'
|
||||
max_seq_len = 1 * 1024
|
||||
lm_config = LMConfig()
|
||||
lm_config.max_seq_len = max_seq_len
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
model, tokenizer, (vision_model, preprocess) = init_model(lm_config)
|
||||
|
||||
model = model.eval()
|
||||
# 推送到huggingface
|
||||
# model.push_to_hub("minimind")
|
||||
# tokenizer.push_to_hub("minimind")
|
||||
|
||||
stream = True
|
||||
|
||||
# prompt = lm_config.image_special_token + '\n这个图片描述的是什么内容？'
|
||||
|
||||
i = 0
|
||||
# 获取图像文件列表
|
||||
image_dir = './dataset/eval_multi_images/'
|
||||
image_files = sorted(os.listdir(image_dir))
|
||||
|
||||
image_count = len(image_files)
|
||||
special_tokens = lm_config.image_special_token * image_count # 根据图片数量重复特殊标记
|
||||
|
||||
# 构建提示
|
||||
prompt = f"{special_tokens}\n这几张图片是否是同一个类别？"
|
||||
|
||||
# 遍历图像文件
|
||||
image_encoders = []
|
||||
for i, image_file in enumerate(image_files):
|
||||
image_path = os.path.join(image_dir, image_file)
|
||||
|
||||
# 打开图像并转换为RGB格式
|
||||
image = Image.open(image_path).convert('RGB')
|
||||
image_process = get_img_process(image, preprocess).to(vision_model.device)
|
||||
|
||||
# 对图像进行编码
|
||||
image_encoder = get_img_embedding(image_process, vision_model).unsqueeze(0)
|
||||
print(image_encoder.shape)
|
||||
image_encoders.append(image_encoder)
|
||||
print(f'[Image]: {image_file}')
|
||||
image_encoders = torch.cat(image_encoders, dim=0).unsqueeze(0)
|
||||
print(image_encoders.shape)
|
||||
|
||||
# 准备问题和消息
|
||||
prompt_ = prompt.replace("\n", "")
|
||||
print(f'[Q]: {prompt_}')
|
||||
|
||||
messages = [{"role": "user", "content": prompt}]
|
||||
|
||||
new_prompt = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)[-(max_seq_len - 1):]
|
||||
|
||||
x = tokenizer(new_prompt).data['input_ids']
|
||||
x = (torch.tensor(x, dtype=torch.long, device=device)[None, ...])
|
||||
|
||||
with torch.no_grad():
|
||||
res_y = model.generate(x, tokenizer.eos_token_id, max_new_tokens=max_seq_len, temperature=temperature,
|
||||
top_k=top_k, stream=stream, image_encoders=image_encoders)
|
||||
|
||||
print('[A]: ', end='')
|
||||
try:
|
||||
y = next(res_y)
|
||||
except StopIteration:
|
||||
print("No answer")
|
||||
exit()
|
||||
|
||||
history_idx = 0
|
||||
while y is not None:
|
||||
answer = tokenizer.decode(y[0].tolist())
|
||||
if answer and answer[-1] == '<EFBFBD>':
|
||||
try:
|
||||
y = next(res_y)
|
||||
except StopIteration:
|
||||
break
|
||||
continue
|
||||
if not len(answer):
|
||||
try:
|
||||
y = next(res_y)
|
||||
except StopIteration:
|
||||
break
|
||||
continue
|
||||
|
||||
print(answer[history_idx:], end='', flush=True)
|
||||
try:
|
||||
y = next(res_y)
|
||||
except StopIteration:
|
||||
break
|
||||
history_idx = len(answer)
|
||||
|
||||
if not stream:
|
||||
break
|
||||
|
||||
print('\n')
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 97 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 109 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 122 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 152 KiB |
|
|
@ -15,7 +15,7 @@ class LMConfig(PretrainedConfig):
|
|||
hidden_dim: int = None,
|
||||
multiple_of: int = 64,
|
||||
norm_eps: float = 1e-5,
|
||||
max_seq_len: int = 200,
|
||||
max_seq_len: int = 512,
|
||||
dropout: float = 0.0,
|
||||
flash_attn: bool = True,
|
||||
image_special_token: str = '<' * 25 + '>' * 25,
|
||||
|
|
|
|||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -180,5 +180,99 @@ class SFTDataset(Dataset):
|
|||
return X_tensor, Y_tensor, loss_mask_tensor, image_encoders
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pass
|
||||
class SFTDataset_multi(Dataset):
|
||||
def __init__(self, json_path, tokenizer, vision_model=None, max_length=1024,
|
||||
prompt_max_len=512,
|
||||
answer_max_len=256,
|
||||
image_special_token='<' * 25 + '>' * 25):
|
||||
|
||||
super().__init__()
|
||||
with open(json_path, 'r', encoding='utf-8') as f:
|
||||
self.data = json.load(f)
|
||||
|
||||
self.max_length = max_length
|
||||
self.prompt_max_len = prompt_max_len
|
||||
self.answer_max_len = answer_max_len
|
||||
self.tokenizer = tokenizer
|
||||
self.vision_model, self.preprocess = vision_model
|
||||
self.padding = 0
|
||||
self.bos_id = self.tokenizer('<s>assistant').data['input_ids']
|
||||
self.dataset_path = './dataset/sft2_images/'
|
||||
self.image_special_token = image_special_token
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def find_sublist_index(self, main_list, sub_list) -> int:
|
||||
last_index = -1
|
||||
for i in range(len(main_list) - len(sub_list) + 1):
|
||||
if main_list[i:i + len(sub_list)] == sub_list:
|
||||
last_index = i
|
||||
return last_index
|
||||
|
||||
def safe_eval(self, s):
|
||||
try:
|
||||
res = eval(s)
|
||||
except Exception as e:
|
||||
return []
|
||||
return res
|
||||
|
||||
def __getitem__(self, index: int):
|
||||
sample = self.data[index]
|
||||
image_names = sample['image'].split(', ')
|
||||
conversation = sample['conversations']
|
||||
|
||||
messages = []
|
||||
for i in range(0, 1):
|
||||
if i + 1 < len(conversation):
|
||||
q = conversation[i]['value'].replace('<image>', self.image_special_token)
|
||||
a = conversation[i + 1]['value']
|
||||
|
||||
if q and a:
|
||||
messages.append({"role": "user", "content": q})
|
||||
messages.append({"role": "assistant", "content": a})
|
||||
|
||||
# print(messages) # [{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nAre the two images below that resemble each other described by the same term? You must choose your answer from the Choice List. Choice_List: True, False.'}, {'role': 'assistant', 'content': 'False'}][{'role': 'user', 'content': '<<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<<<<<<>>>>>>>>>>>>>>>>>>>>>>>>>\nDo these four pictures fall into the same category? You must choose your answer from the Choice List. Choice_List: True, False.'}, {'role': 'assistant', 'content': 'False'}]
|
||||
new_prompt = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True
|
||||
)
|
||||
input_id = self.tokenizer(new_prompt).data['input_ids'][:self.max_length]
|
||||
# print(len(input_id)) # 165 or 263 or 259
|
||||
|
||||
question_length = self.find_sublist_index(input_id, self.bos_id) + len(self.bos_id)
|
||||
padding_len = self.max_length - len(input_id)
|
||||
input_id = input_id + [self.padding] * padding_len
|
||||
mask_len = len(input_id) - question_length - padding_len
|
||||
loss_mask = [0] * question_length + [1] * (mask_len) + [0] * padding_len
|
||||
|
||||
input_id = np.array(input_id)
|
||||
|
||||
X = np.array(input_id[:-1]).astype(np.int64)
|
||||
Y = np.array(input_id[1:]).astype(np.int64)
|
||||
loss_mask = np.array(loss_mask[1:]).astype(np.int64)
|
||||
|
||||
X_tensor = torch.from_numpy(X)
|
||||
Y_tensor = torch.from_numpy(Y)
|
||||
loss_mask_tensor = torch.from_numpy(loss_mask)
|
||||
|
||||
# 读取多张图像
|
||||
image_encoders = []
|
||||
for image_name in image_names:
|
||||
image = Image.open(f'{self.dataset_path}{image_name.strip()}') # 去掉可能的空格
|
||||
image_encoders.append(get_img_process(image, self.preprocess)['pixel_values'])
|
||||
|
||||
# 确定目标形状
|
||||
max_images = 4 # 根据你的需求设置 一次性最大输入几张图片
|
||||
target_shape = (max_images, 3, 224, 224)
|
||||
|
||||
# 创建填充张量
|
||||
padded_image_encoders = torch.zeros(target_shape, dtype=torch.float32)
|
||||
|
||||
# 填充图像编码
|
||||
for i, img_enc in enumerate(image_encoders):
|
||||
if i < max_images:
|
||||
padded_image_encoders[i] = img_enc.squeeze(0)
|
||||
|
||||
return X_tensor, Y_tensor, loss_mask_tensor, padded_image_encoders
|
||||
|
|
@ -372,29 +372,38 @@ class Transformer(PreTrainedModel):
|
|||
# 查找token中<image>片段的索引，为了替换做准备
|
||||
def find_indices(tokens, image_ids):
|
||||
image_ids_tensor = torch.tensor(image_ids).to(tokens.device)
|
||||
return [
|
||||
[i, i + len(image_ids) - 1]
|
||||
for batch_idx in range(tokens.size(0))
|
||||
for i in range(tokens.size(1) - len(image_ids) + 1)
|
||||
if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor)
|
||||
] or None
|
||||
indices = []
|
||||
|
||||
image_indices = find_indices(tokens, self.image_ids)
|
||||
for batch_idx in range(tokens.size(0)):
|
||||
for i in range(tokens.size(1) - len(image_ids) + 1):
|
||||
if torch.equal(tokens[batch_idx, i:i + len(image_ids)], image_ids_tensor):
|
||||
indices.append([batch_idx, i, i + len(image_ids) - 1]) # 返回batch_idx和开始结束索引
|
||||
|
||||
return indices if indices else None
|
||||
|
||||
image_indices = find_indices(tokens, self.image_ids) # [0, 4, 53], [0, 54, 103], [0, 104, 153], [0, 154, 203] or [1, 4, 53], [1, 54, 103]
|
||||
|
||||
# 如果此时有图像编码
|
||||
if image_encoders is not None:
|
||||
vision_proj = self.vision_proj(image_encoders)
|
||||
vision_proj = self.vision_proj(image_encoders) # torch.Size([32, 4, 50, 512])
|
||||
if image_indices is not None:
|
||||
# 创建一个新的张量来存储拼接后的结果
|
||||
new_h = []
|
||||
|
||||
for i in range(h.size(0)):
|
||||
before = h[i, :image_indices[i][0], :]
|
||||
after = h[i, image_indices[i][1] + 1:, :]
|
||||
current_h = h[i]
|
||||
img_idx = 0
|
||||
for _, (batch_idx, start_idx, end_idx) in enumerate(image_indices):
|
||||
if batch_idx == i:
|
||||
# 插入vision_proj特征
|
||||
before = current_h[:start_idx, :]
|
||||
after = current_h[end_idx + 1:, :]
|
||||
# 拼接 before, vision_proj, after
|
||||
new_h_i = torch.cat((before, vision_proj[i], after), dim=0)[:seqlen]
|
||||
new_h.append(new_h_i)
|
||||
current_h = torch.cat((before, vision_proj[i][img_idx], after), dim=0)[:seqlen]
|
||||
img_idx += 1
|
||||
new_h.append(current_h)
|
||||
# 将所有拼接后的结果堆叠起来
|
||||
new_h = torch.stack(new_h, dim=0)
|
||||
new_h = torch.stack(new_h, dim=0) # torch.Size([32, 511, 512])
|
||||
return new_h
|
||||
|
||||
return h
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from transformers import CLIPProcessor, CLIPModel
|
|||
from PIL import Image
|
||||
import requests
|
||||
import torch
|
||||
import transformers
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
|
|
@ -31,7 +32,10 @@ def get_img_embedding(batch_encoding, vision_model):
|
|||
embeddings.append(output.last_hidden_state)
|
||||
|
||||
# 从 BatchEncoding 中提取图像张量
|
||||
image_tensor = batch_encoding['pixel_values'] # 假设图像张量在 'pixel_values' 键下
|
||||
if isinstance(batch_encoding, transformers.tokenization_utils_base.BatchEncoding):
|
||||
image_tensor = batch_encoding['pixel_values']
|
||||
else:
|
||||
image_tensor = batch_encoding # torch.Size([32, 4, 3, 224, 224])
|
||||
|
||||
# 如果图像张量的形状是5维，则无需添加额外维度
|
||||
if len(image_tensor.shape) == 4:
|
||||
|
|
@ -54,5 +58,5 @@ def get_img_embedding(batch_encoding, vision_model):
|
|||
hook.remove()
|
||||
|
||||
# 拼接所有特征向量成为一个张量
|
||||
all_embeddings = torch.stack(embeddings, dim=0).squeeze()
|
||||
all_embeddings = torch.stack(embeddings, dim=0).squeeze() # torch.Size([32, 4, 50, 768])
|
||||
return all_embeddings
|
||||
|
|
|
|||
|
|
@ -0,0 +1,53 @@
|
|||
import json
|
||||
|
||||
# 加载现有JSON数据
|
||||
with open('full.json', 'r') as f:
|
||||
data = json.load(f)
|
||||
|
||||
# 创建新的格式
|
||||
new_data = []
|
||||
|
||||
# 定义选择列表的翻译字典
|
||||
translation_dict = {
|
||||
"True": "是的",
|
||||
"False": "不是",
|
||||
"painted": "涂鸦",
|
||||
"ripe": "成熟"
|
||||
}
|
||||
|
||||
for item in data['data']:
|
||||
sample_id = str(item['sample_id']).zfill(12) # 格式化为12位数字
|
||||
images = [f"<image>" for _ in item['task_instance']['images_path']]
|
||||
|
||||
# 直接从task_instance中获取choice_list
|
||||
choice_list = item['task_instance']['choice_list']
|
||||
|
||||
# 判断是否需要翻译
|
||||
if all(choice in translation_dict for choice in choice_list):
|
||||
choice_list = [translation_dict[choice] for choice in choice_list]
|
||||
response = translation_dict[item['response']]
|
||||
else:
|
||||
response = item['response']
|
||||
|
||||
# 构造新的对话格式
|
||||
conversation = [
|
||||
{
|
||||
"from": "human",
|
||||
"value": f"{''.join(images)}\n{data['metadata']['task_instruction'][item['task_instruction_id']]} 选择列表: {', '.join(choice_list)}."
|
||||
},
|
||||
{
|
||||
"from": "gpt",
|
||||
"value": response
|
||||
}
|
||||
]
|
||||
|
||||
# 添加到新数据中
|
||||
new_data.append({
|
||||
"id": sample_id,
|
||||
"image": ', '.join(item['task_instance']['images_path']),
|
||||
"conversations": conversation
|
||||
})
|
||||
|
||||
# 保存为新的JSON格式
|
||||
with open('output.json', 'w') as f:
|
||||
json.dump(new_data, f, ensure_ascii=False, indent=2)
|
||||
Loading…
Reference in New Issue