与vlm统一流程
forXuyx
parent
c843572b59
commit
cfac34ca6a
|
|
@ -103,21 +103,18 @@ class VLMDataset(Dataset):
|
|||
return X, Y, loss_mask, image_tensors
|
||||
|
||||
class T2IDataset(Dataset):
|
||||
def __init__(self, jsonl_path, images_path, tokenizer, max_length=512,
|
||||
def __init__(self, jsonl_path, images_path, tokenizer, max_length=512, img_pre_process=False,
|
||||
image_special_token='@' * 256):
|
||||
|
||||
super().__init__()
|
||||
self.samples = self.load_data(jsonl_path)
|
||||
self.images_path = images_path
|
||||
|
||||
self.img_pre_process = img_pre_process # 是否使用提前处理的图片
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.max_length = max_length
|
||||
self.image_size = 256
|
||||
self.transform = transforms.Compose([
|
||||
transforms.Lambda(lambda pil_image: center_crop_arr(pil_image, self.image_size)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
|
||||
])
|
||||
self.image_token = image_special_token
|
||||
self.bos_id = tokenizer('<s>assistant\n', add_special_tokens=False).input_ids
|
||||
self.eos_id = tokenizer('</s>\n', add_special_tokens=False).input_ids
|
||||
|
|
@ -174,12 +171,26 @@ class T2IDataset(Dataset):
|
|||
Y = torch.tensor(input_ids[1:], dtype=torch.long)
|
||||
loss_mask = torch.tensor(loss_mask[1:], dtype=torch.long)
|
||||
|
||||
image_tensors = []
|
||||
for image_name in image_paths.split(','):
|
||||
image_name = image_name.strip()
|
||||
image = Image.open(f'{self.images_path}/{image_name}')
|
||||
image_tensor = self.transform(image)
|
||||
image_tensors.append(image_tensor)
|
||||
image_tensors = torch.stack(image_tensors, dim=0)
|
||||
if self.img_pre_process:
|
||||
image_emb_path = image_paths.replace('.jpg', '_emb.npy')
|
||||
image_emb = np.load(f'{self.images_path}/{image_emb_path}')
|
||||
image_emb = torch.tensor(image_emb, dtype=torch.float32)
|
||||
# 加载预处理的图像token
|
||||
image_token_path = image_paths.replace('.jpg', '_token.npy')
|
||||
image_token = np.load(f'{self.images_path}/{image_token_path}')
|
||||
image_token = torch.tensor(image_token, dtype=torch.long)
|
||||
return X, Y, loss_mask, (image_emb, image_token)
|
||||
else:
|
||||
image_tensors = []
|
||||
for image_name in image_paths.split(','):
|
||||
image_name = image_name.strip()
|
||||
image = Image.open(f'{self.images_path}/{image_name}').convert("RGB")
|
||||
image = center_crop_arr(image, self.image_size)
|
||||
image = np.array(image) / 255.
|
||||
image = 2.0 * image - 1.0
|
||||
image = torch.tensor(image, dtype=torch.float32)
|
||||
image_tensor = torch.einsum('hwc->chw', image)
|
||||
image_tensors.append(image_tensor)
|
||||
image_tensors = torch.stack(image_tensors, dim=0)
|
||||
|
||||
return X, Y, loss_mask, image_tensors
|
||||
return X, Y, loss_mask, image_tensors
|
||||
|
|
|
|||
|
|
@ -3,13 +3,29 @@ from .model import *
|
|||
from typing import Optional, Tuple, List
|
||||
from torch import nn
|
||||
import warnings
|
||||
from model.model_vlm import MiniMindVLM, VisionProj
|
||||
from model.model_vlm import MiniMindVLM
|
||||
from model.model_vq import VQ_models
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
warnings.filterwarnings('ignore')
|
||||
|
||||
class VisionProj(nn.Module):
|
||||
def __init__(self, ve_dim=768, lm_dim=512, hidden_dim=256):
|
||||
super().__init__()
|
||||
self.ve_dim = ve_dim
|
||||
self.lm_dim = lm_dim
|
||||
self.hidden_dim = hidden_dim
|
||||
self.vision_proj = nn.Sequential(
|
||||
nn.Linear(self.ve_dim, self.hidden_dim),
|
||||
nn.ReLU(),
|
||||
nn.Linear(self.hidden_dim, self.lm_dim)
|
||||
)
|
||||
|
||||
def forward(self, image_encoders):
|
||||
vision_proj = self.vision_proj(image_encoders)
|
||||
return vision_proj
|
||||
|
||||
|
||||
# 继承自语言模型
|
||||
class MiniMindT2I(MiniMindVLM):
|
||||
|
|
@ -20,14 +36,12 @@ class MiniMindT2I(MiniMindVLM):
|
|||
if not params: params = VLMConfig()
|
||||
self.params = params
|
||||
self.image_tokenizer = self.__class__.get_image_tokenizer()
|
||||
# 这里重写是为了将output映射到codebook_size
|
||||
self.output = nn.Linear(params.dim, 16384, bias=False)
|
||||
self.vision_proj = VisionProj(ve_dim=8, lm_dim=params.dim)
|
||||
|
||||
@staticmethod
|
||||
def get_image_tokenizer(model_path="./model/minimind_img_tokenizer/vq_ds16_t2i.pt"):
|
||||
def get_image_tokenizer(model_path="./model/minimind_img_tokenizer/minimind_img_tokenizer.pt"):
|
||||
model = VQ_models['VQ-16'](
|
||||
codebook_size=16384,
|
||||
codebook_size=6400,
|
||||
codebook_embed_dim=8)
|
||||
|
||||
# 加载模型
|
||||
|
|
@ -105,23 +119,31 @@ class MiniMindT2I(MiniMindVLM):
|
|||
h = self.tok_embeddings(input_ids)
|
||||
|
||||
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
|
||||
stack_dim = 1 if bs > 1 else 0
|
||||
# 获取图片的embedding
|
||||
vision_tensors = torch.stack([
|
||||
MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
|
||||
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])
|
||||
# 获取图片的token
|
||||
vision_tokens = torch.stack([
|
||||
MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
|
||||
for i in range(num)
|
||||
], dim=stack_dim)
|
||||
# 替换target中的'@......'为图片的token
|
||||
target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
|
||||
if isinstance(pixel_tensors, torch.Tensor):
|
||||
pixel_tensors = pixel_tensors.to(h.device)
|
||||
if len(pixel_tensors.shape) == 6:
|
||||
pixel_tensors = pixel_tensors.squeeze(2)
|
||||
bs, num, c, im_h, im_w = pixel_tensors.shape
|
||||
stack_dim = 1 if bs > 1 else 0
|
||||
# 获取图片的embedding
|
||||
vision_tensors = torch.stack([
|
||||
MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
|
||||
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])
|
||||
# 获取图片的token
|
||||
vision_tokens = torch.stack([
|
||||
MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
|
||||
for i in range(num)
|
||||
], dim=stack_dim)
|
||||
# 替换target中的'@......'为图片的token
|
||||
target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
|
||||
else:
|
||||
vision_tensors, vision_tokens = pixel_tensors
|
||||
vision_tensors = vision_tensors.to(h.device)
|
||||
vision_tokens = vision_tokens.to(h.device)
|
||||
h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
|
||||
target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
|
||||
|
||||
pos_cis = self.pos_cis[start_pos:start_pos + input_ids.shape[1]]
|
||||
past_kvs = []
|
||||
|
|
|
|||
|
|
@ -39,13 +39,14 @@ def train_epoch(epoch, wandb):
|
|||
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_tensors = pixel_tensors
|
||||
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, Y, pixel_tensors=pixel_tensors)
|
||||
Y = res.target_ids
|
||||
loss = loss_fct(
|
||||
res.logits.view(-1, res.logits.size(-1)),
|
||||
res.target_ids.view(-1)
|
||||
|
|
@ -147,7 +148,7 @@ if __name__ == "__main__":
|
|||
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_t2i_data.jsonl")
|
||||
parser.add_argument("--images_path", type=str, default="./dataset/pretrain_images")
|
||||
parser.add_argument("--images_path", type=str, default="./dataset/pretrain_t2i_code")
|
||||
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)
|
||||
|
|
@ -159,6 +160,7 @@ if __name__ == "__main__":
|
|||
parser.add_argument('--n_layers', default=8, type=int)
|
||||
parser.add_argument('--max_seq_len', default=640, type=int)
|
||||
parser.add_argument('--use_moe', default=False, type=bool)
|
||||
parser.add_argument('--img_process', type=bool, default=True)
|
||||
args = parser.parse_args()
|
||||
|
||||
model_config = VLMConfig(dim=args.dim, n_layers=args.n_layers, max_seq_len=args.max_seq_len,
|
||||
|
|
@ -189,9 +191,8 @@ if __name__ == "__main__":
|
|||
|
||||
model, tokenizer = init_model(model_config)
|
||||
|
||||
train_ds = T2IDataset(args.data_path, args.images_path, tokenizer,
|
||||
image_special_token=model_config.image_special_token,
|
||||
max_length=max_seq_len)
|
||||
train_ds = T2IDataset(args.data_path, args.images_path, tokenizer, max_seq_len, args.img_process,
|
||||
image_special_token=model_config.image_special_token)
|
||||
train_sampler = DistributedSampler(train_ds) if ddp else None
|
||||
train_loader = DataLoader(
|
||||
train_ds,
|
||||
|
|
|
|||
Loading…
Reference in New Issue