更新

eval_t2i.py

@@ -0,0 +1,93 @@
+import argparse
+import os
+import random
+import numpy as np
+import torch
+import warnings
+import torch.nn.functional as F
+from PIL import Image
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from model.model_t2i import MiniMindT2I
+from model.VLMConfig import VLMConfig
+from transformers import logging as hf_logging
+
+hf_logging.set_verbosity_error()
+
+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, device):
+    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    moe_path = '_moe' if args.use_moe else ''
+    ckp = f'./{args.out_dir}/sft_t2i_{args.dim}{moe_path}.pth'
+    model = MiniMindT2I(lm_config)
+    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)
+
+    print(f'T2I参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
+
+    return model.eval().to(device), tokenizer
+
+
+def setup_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Chat with MiniMind")
+    parser.add_argument('--lora_name', default='None', type=str)
+    parser.add_argument('--out_dir', default='out', type=str)
+    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)
+    parser.add_argument('--max_seq_len', default=640, type=int)
+    parser.add_argument('--use_moe', default=False, type=bool)
+    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)
+
+    model, tokenizer = init_model(lm_config, args.device)
+
+
+    def chat_with_vlm(prompt):
+        messages = [{"role": "user", "content": prompt}]
+
+        new_prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )[-args.max_seq_len + 1:]
+
+        with torch.no_grad():
+            x = torch.tensor(tokenizer(new_prompt)['input_ids'], device=args.device).unsqueeze(0)
+            outputs = model.generate(
+                x,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                pad_token_id=tokenizer.pad_token_id,
+            )
+            image = model.image_tokenizer.decode_code(outputs, (1, 8, 16, 16))
+            image = F.interpolate(image, size=[256, 256], mode='bicubic').permute(0, 2, 3, 1)[0]
+            image = torch.clamp(127.5 * image + 128.0, 0, 255).to("cpu", dtype=torch.uint8).numpy()
+            Image.fromarray(image).save('output.jpg')
+            print('🤖️: 图片保存至output.jpg')
+            print('\n')
+
+
+    image_dir = './dataset/eval_images/'
+    prompt = f"一个年轻人准备踢足球"
+    chat_with_vlm(prompt)

model/model_t2i.py

@@ -3,28 +3,28 @@ from .model import *
 from typing import Optional, Tuple, List
 from torch import nn
 import warnings
-from model.model_vlm import MiniMindVLM
+from model.model_vlm import MiniMindVLM, VisionProj
 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):
+# class VisionProj(nn.Module):
-        vision_proj = self.vision_proj(image_encoders)
+#     def __init__(self, ve_dim=768, lm_dim=512):
-        return vision_proj
+#         super().__init__()
+#         self.ve_dim = ve_dim
+#         self.lm_dim = lm_dim
+#         self.vision_proj = nn.Sequential(
+#             nn.Linear(self.ve_dim, self.lm_dim),
+#             nn.ReLU(),
+#             nn.Linear(self.lm_dim, self.lm_dim)
+#         )
+
+#     def forward(self, image_encoders):
+#         vision_proj = self.vision_proj(image_encoders)
+#         return vision_proj
 
 
 # 继承自语言模型
@@ -36,6 +36,7 @@ class MiniMindT2I(MiniMindVLM):
         if not params: params = VLMConfig()
         self.params = params
         self.image_tokenizer = self.__class__.get_image_tokenizer()
+        # self.output = nn.Linear(params.dim, 16384)
         self.vision_proj = VisionProj(ve_dim=8, lm_dim=params.dim)
 
     @staticmethod
@@ -116,34 +117,39 @@ class MiniMindT2I(MiniMindVLM):
                 **args):
         start_pos = args.get('start_pos', 0)
         pixel_tensors = args.get('pixel_tensors', None)
-        h = self.tok_embeddings(input_ids)
+        is_image = args.get('is_image', False)
+        if is_image:
+            input_ids = self.image_tokenizer.quantize.get_codebook_entry(input_ids)
+            h = self.vision_proj(input_ids)
+        else:
+            h = self.tok_embeddings(input_ids)
 
-        if pixel_tensors is not None and start_pos == 0:
+            if pixel_tensors is not None and start_pos == 0:
-            if isinstance(pixel_tensors, torch.Tensor):
+                if isinstance(pixel_tensors, torch.Tensor):
-                pixel_tensors = pixel_tensors.to(h.device)
+                    pixel_tensors = pixel_tensors.to(h.device)
-                if len(pixel_tensors.shape) == 6:
+                    if len(pixel_tensors.shape) == 6:
-                    pixel_tensors = pixel_tensors.squeeze(2)
+                        pixel_tensors = pixel_tensors.squeeze(2)
-                bs, num, c, im_h, im_w = pixel_tensors.shape
+                    bs, num, c, im_h, im_w = pixel_tensors.shape
-                stack_dim = 1 if bs > 1 else 0
+                    stack_dim = 1 if bs > 1 else 0
-                # 获取图片的embedding
+                    # 获取图片的embedding
-                vision_tensors = torch.stack([
+                    vision_tensors = torch.stack([
-                    MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
+                        MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[0]
-                    for i in range(num)
+                        for i in range(num)
-                ], dim=stack_dim)
+                    ], dim=stack_dim)
-                h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
+                    h = self.count_vision_proj(tokens=input_ids, h=h, vision_tensors=vision_tensors, seqlen=input_ids.shape[1])
-                # 获取图片的token
+                    # 获取图片的token
-                vision_tokens = torch.stack([
+                    vision_tokens = torch.stack([
-                    MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
+                        MiniMindT2I.get_image_embeddings(pixel_tensors[:, i, :, :, :], self.image_tokenizer)[1]
-                    for i in range(num)
+                        for i in range(num)
-                ], dim=stack_dim)
+                    ], dim=stack_dim)
-                # 替换target中的'@......'为图片的token
+                    # 替换target中的'@......'为图片的token
-                target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
+                    target_ids = self.count_token_replace(tokens=target_ids, vision_token=vision_tokens, seqlen=target_ids.shape[1])
-            else:
+                else:
-                vision_tensors, vision_tokens = pixel_tensors
+                    vision_tensors, vision_tokens = pixel_tensors
-                vision_tensors = vision_tensors.to(h.device)
+                    vision_tensors = vision_tensors.to(h.device)
-                vision_tokens = vision_tokens.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])
+                    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])
+                    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 = []
@@ -163,3 +169,31 @@ class MiniMindT2I(MiniMindVLM):
         self.OUT.__setitem__('past_key_values', past_kvs)
         self.OUT.__setitem__('target_ids', target_ids) # 添加target_ids
         return self.OUT
+
+    @torch.inference_mode()
+    def generate(self, input_ids, img_token_num=256, temperature=0.75, top_p=0.90, rp=1., use_cache=True, pad_token_id=0, **args):
+        idx = 0
+        input_ids = input_ids[input_ids != pad_token_id].unsqueeze(0)
+        start, first_seq, past_kvs = input_ids.shape[1], True, None
+        while idx < img_token_num:
+            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, is_image=True, **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)
+            idx += 1
+        return input_ids[:, -img_token_num:]

train_pretrain_t2i.py

@@ -87,7 +87,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_t2i_{model_config.dim}{moe_path}.pth'
             if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                 state_dict = model.module.state_dict()
             else:

train_sft_t2i.py

@@ -0,0 +1,205 @@
+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, nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, DistributedSampler
+from transformers import AutoTokenizer, AutoModel
+from model.model_t2i import MiniMindT2I
+from model.VLMConfig import VLMConfig
+from model.dataset import T2IDataset
+
+warnings.filterwarnings('ignore')
+
+
+def Logger(content):
+    if not ddp or dist.get_rank() == 0:
+        print(content)
+
+
+def get_lr(current_step, total_steps, lr):
+    return lr / 10 + 0.5 * lr * (1 + math.cos(math.pi * current_step / total_steps))
+
+
+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):
+        X = X.to(args.device)
+        Y = Y.to(args.device)
+        loss_mask = loss_mask.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)
+            ).view(Y.size())
+
+            loss = (loss * loss_mask).sum() / loss_mask.sum()
+            loss += res.aux_loss
+            loss = loss / args.accumulation_steps
+
+        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 + 1,
+                    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 model_config.use_moe else ''
+            ckp = f'{args.save_dir}/sft_t2i_{model_config.dim}{moe_path}.pth'
+            if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+                state_dict = model.module.state_dict()
+            else:
+                state_dict = model.state_dict()
+            clean_state_dict = {
+                key: value for key, value in state_dict.items() if not key.startswith('vision_encoder.')
+            }
+            torch.save(clean_state_dict, ckp)
+            model.train()
+
+
+def init_model(model_config: VLMConfig):
+    tokenizer = AutoTokenizer.from_pretrained('./model/minimind_tokenizer')
+    moe_path = '_moe' if model_config.use_moe else ''
+    # 加载纯语言模型权重
+    ckp = f'./out/pretrain_t2i_{model_config.dim}{moe_path}.pth'
+    model = MiniMindT2I(model_config)
+    state_dict = torch.load(ckp, map_location=args.device)
+    model.load_state_dict(state_dict, strict=False)
+
+    Logger(f'T2I可训练参数量：{sum(p.numel() for p in model.parameters() if p.requires_grad) / 1e6:.3f} 百万')
+
+    return model.to(args.device), tokenizer
+
+
+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 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=16)
+    parser.add_argument("--learning_rate", type=float, default=4e-4)
+    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/pretrain_t2i_data.jsonl")
+    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)
+    parser.add_argument("--warmup_iters", type=int, default=0)
+    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('--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,
+                             image_special_token = '@' * 256, image_ids = [34] * 256)
+    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)
+    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 Pretrain-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 = init_model(model_config)
+
+    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,
+        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.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr=args.learning_rate)
+
+    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)