update pretrain and finetune codes of llama_adapter_v2_multimodal

Author: shilinyan99

llama_adapter_v2_multimodal/data/dataset.py

@@ -0,0 +1,157 @@
+import torch
+import yaml
+from torch.utils.data import Dataset
+from PIL import Image
+import json
+import llama.utils
+from llama import Tokenizer
+import copy
+import torchvision.transforms as transforms
+import pandas as pd
+import random
+import cv2
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+PROMPT_DICT = {
+    "prompt_input": (
+        "Below is an instruction that describes a task, paired with an input that provides further context. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
+    ),
+    "prompt_no_input": (
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+
+# create data
+transform_train = transforms.Compose([
+    transforms.RandomResizedCrop(size=(224, 224), scale=(0.9, 1.0), ratio=(0.75, 1.3333), interpolation=BICUBIC,
+                                 antialias=None),  # 3 is bicubic
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711])])
+
+class FinetuneDataset(Dataset):
+    def __init__(self, config_path, transform, max_words=30, tokenizer_path=None):
+        print(f"read dataset config from {config_path}")
+        with open(config_path, 'r') as f:
+            self.config = yaml.load(f, Loader=yaml.FullLoader)
+        print("DATASET CONFIG:")
+        print(self.config)
+        ann = []
+        for meta_path in self.config['META']:
+            meta_l = json.load(open(meta_path))
+            print(f"{meta_path}: len {len(meta_l)}")
+            ann += meta_l
+        self.ann = ann
+        print(f"total length: {len(self)}")
+        self.transform = transform
+        self.max_words = max_words
+        self.tokenizer = Tokenizer(model_path=tokenizer_path)
+
+    def __len__(self):
+        return len(self.ann)
+
+    def __getitem__(self, index):
+        data_item = self.ann[index]
+        if 'image' in data_item.keys():
+            filename = data_item['image']
+            question = data_item['conversations'][0]['value']
+            answer = data_item['conversations'][1]['value']
+     
+            image = cv2.imread(filename)
+            image = Image.fromarray(image)
+            image = self.transform(image)
+            format_instruction = question
+            format_input = None
+        else:
+            image = torch.zeros(3, 224, 224)
+            format_instruction = data_item['instruction'],
+            format_input = data_item['input']
+            answer = data_item['output']
+        input1 = llama.utils.format_prompt(format_instruction, format_input)
+        input2 = input1 + answer
+        input1 = torch.tensor(self.tokenizer.encode(input1, bos=True, eos=False), dtype=torch.int64)
+        input2 = torch.tensor(self.tokenizer.encode(input2, bos=True, eos=True), dtype=torch.int64)
+        padding = self.max_words - input2.shape[0]
+        if padding > 0:
+            input2 = torch.cat((input2, torch.zeros(padding, dtype=torch.int64) - 1))
+        elif padding < 0:
+            input2 = input2[:self.max_words]
+        labels = copy.deepcopy(input2)
+        labels[:len(input1)] = -1
+        input2_mask = input2.ge(0)
+        label_mask = labels.ge(0)
+        input2[~input2_mask] = 0
+        labels[~label_mask] = 0
+        input2_mask = input2_mask.float()
+        label_mask = label_mask.float()
+        return input2, labels, input2_mask, image
+
+
+class PretrainDataset(Dataset):
+    def __init__(self, config_path, transform, max_words=30, tokenizer_path=None):
+        print(f"read dataset config from {config_path}")
+        with open(config_path, 'r') as f:
+            self.config = yaml.load(f, Loader=yaml.FullLoader)
+        print("DATASET CONFIG:")
+        print(self.config)
+        images, captions = [], []
+        for meta_path in self.config['META']:
+            images_this_meta, captions_this_meta = [], []
+            for chunk in pd.read_csv(meta_path, sep='\t', lineterminator='\n', chunksize=10 ** 6):
+                images_this_meta.extend(chunk['url'].tolist())
+                captions_this_meta.extend(chunk['caption'].tolist())
+            print(f"{meta_path}: len {len(images_this_meta)}")
+            images.extend(images_this_meta)
+            captions.extend(captions_this_meta)
+
+        self.data_list = []
+        for x, y in zip(images, captions):
+            self.data_list.append({'url': x, 'caption': y})
+        print(f"total length: {len(self)}")
+        self.transform = transform
+        self.max_words = max_words
+        self.tokenizer = Tokenizer(model_path=tokenizer_path)
+
+    def __len__(self):
+        return len(self.data_list)
+
+    def __getitem__(self, index):
+        sample = self.data_list[index]
+        image_path, caption = sample['url'], sample['caption']
+        if isinstance(caption, list):
+            caption = random.choice(caption)
+        caption = str(caption)
+
+        image = cv2.imread(image_path)
+        image = Image.fromarray(image)
+        image = self.transform(image)
+
+        format_instruction = "Generate caption of this image"
+        input1 = llama.utils.format_prompt(format_instruction, None)
+        input2 = input1 + caption
+
+        input1 = torch.tensor(self.tokenizer.encode(input1, bos=True, eos=False), dtype=torch.int64)
+        input2 = torch.tensor(self.tokenizer.encode(input2, bos=True, eos=True), dtype=torch.int64)
+        padding = self.max_words - input2.shape[0]
+        if padding > 0:
+            input2 = torch.cat((input2, torch.zeros(padding, dtype=torch.int64) - 1))
+        elif padding < 0:
+            input2 = input2[:self.max_words]
+        labels = copy.deepcopy(input2)
+        labels[:len(input1)] = -1
+        input2_mask = input2.ge(0)
+        label_mask = labels.ge(0)
+        input2[~input2_mask] = 0
+        labels[~label_mask] = 0
+        input2_mask = input2_mask.float()
+        label_mask = label_mask.float()
+        return input2, labels, input2_mask, image

llama_adapter_v2_multimodal/demo.py

@@ -8,9 +8,10 @@
 llama_dir = "/path/to/LLaMA/"
 
 model, preprocess = llama.load("BIAS-7B", llama_dir, device)
+model.eval()
 
 prompt = llama.format_prompt('Please introduce this painting.')
-img = Image.fromarray(cv2.imread("../docs/logo_v1.png"))
+img = Image.fromarray(cv2.imread("./docs/logo_v1.png"))
 img = preprocess(img).unsqueeze(0).to(device)
 
 result = model.generate(img, [prompt])[0]

llama_adapter_v2_multimodal/docs/train.md

@@ -0,0 +1,67 @@
+The training process of LLaMA-Adapter V2 consists of the pre-training and fine-tuning phases. 
+
+## Pre-training
+### Data
+* We use multiple datasets with **image-text pairs** for pre-training. The texts are English-only.
+
+* For each dataset, the meta file should be organized in the `.csv` format as following:
+
+  ```
+  url		caption
+  /path/to/image1		caption1
+  /path/to/image2		caption2
+  ...
+  ```
+
+  Alternatively, you may modify the [`PretrainDataset`](/data/dataset.py) implementation to adapt to your own meta file format.
+
+* Write a `.yaml` config file to specify the datasets for pre-training:
+  ```
+  META:
+    - '/path/to/cc3m.csv'
+    - '/path/to/coco.csv'
+    ...
+  ```
+
+### Start pre-training
+
+We are now ready to start pre-training (please make sure that the original LLaMA / Open-Chinese-LLaMA weights are available in `/path/to/llama_model_weights`). 
+
+```bash
+. exps/pretrain.sh /path/to/llama_model_weights /path/to/pretrain-data-config.yaml /output/path
+```
+
+
+
+## Fine-tuning
+
+### Data
+
+* We fine-tune LLaMA-Adapter V2 on text-only as well as image-text instruction following datasets.
+
+* The following lists the datasets we use for training our release weights:
+
+  | Name                     | Link                                                         |
+  | ------------------------ | ------------------------------------------------------------ |
+  | alpaca_gpt4_data.json    | [File Link](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM/blob/main/data/alpaca_gpt4_data.json) |
+  | alpaca_gpt4_data_zh.json | [File Link](https://github.com/Instruction-Tuning-with-GPT-4/GPT-4-LLM/blob/main/data/alpaca_gpt4_data_zh.json) |
+  | llava_instruct_150k.json | [File Link](https://huggingface.co/datasets/liuhaotian/LLaVA-Instruct-150K/raw/main/llava_instruct_150k.json) |
+  | alpaca_data_zh_51k.json  | [File Link](https://github.com/ymcui/Chinese-LLaMA-Alpaca/blob/main/data/alpaca_data_zh_51k.json) |
+
+* Similar to pre-training, write a `.yaml` config file to specify the datasets for fine-tuning:
+
+  ```
+  META:
+    - '/path/to/alpaca_gpt4_data.json'
+    - '/path/to/alpaca_gpt4_data_zh.json'
+    ...
+  ```
+
+### Start fine-tuning
+
+```bash
+. exps/finetune.sh \
+ /path/to/llama_model_weights /path/to/pre-trained/checkopint.pth \
+ /path/to/finetune-data-config.yaml /output/path
+```
+

llama_adapter_v2_multimodal/engine_finetune.py

@@ -0,0 +1,77 @@
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+
+from llama import LLaMA_adapter
+
+def train_one_epoch(model: LLaMA_adapter,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    log_writer=None,
+                    args=None):
+    model.train(True)
+    # model.module.set_default_trainability()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 10
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+    for data_iter_step, (examples, labels, example_mask, imgs) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        imgs = imgs.to(device, non_blocking=True)
+        with torch.cuda.amp.autocast():
+             c_loss, m_loss = model(examples, labels, imgs)
+        loss = c_loss  + m_loss * 0
+        loss_value = loss.item()
+        c_loss_value = c_loss.item()
+        m_loss_value = m_loss
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(closs=c_loss_value)
+        metric_logger.update(mloss=m_loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        c_loss_value_reduce = misc.all_reduce_mean(c_loss_value)
+        m_loss_value_reduce = misc.all_reduce_mean(m_loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('c_train_loss', c_loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('m_train_loss', m_loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}

llama_adapter_v2_multimodal/engine_pretrain.py

@@ -0,0 +1,77 @@
+import math
+import sys
+from typing import Iterable
+
+import torch
+
+import util.misc as misc
+import util.lr_sched as lr_sched
+
+from llama import LLaMA_adapter
+
+def train_one_epoch(model: LLaMA_adapter,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, loss_scaler,
+                    log_writer=None,
+                    args=None):
+    model.train(True)
+    # model.module.set_default_trainability()
+
+    metric_logger = misc.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', misc.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 10
+
+    accum_iter = args.accum_iter
+
+    optimizer.zero_grad()
+
+    if log_writer is not None:
+        print('log_dir: {}'.format(log_writer.log_dir))
+    for data_iter_step, (examples, labels, example_mask, imgs) in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
+        # we use a per iteration (instead of per epoch) lr scheduler
+        if data_iter_step % accum_iter == 0:
+            lr_sched.adjust_learning_rate(optimizer, data_iter_step / len(data_loader) + epoch, args)
+
+        imgs = imgs.to(device, non_blocking=True)
+        with torch.cuda.amp.autocast():
+             c_loss, m_loss = model(examples, labels, imgs)
+        loss = c_loss  + m_loss * 0
+        loss_value = loss.item()
+        c_loss_value = c_loss.item()
+        m_loss_value = m_loss
+        if not math.isfinite(loss_value):
+            print("Loss is {}, stopping training".format(loss_value))
+            sys.exit(1)
+
+        loss /= accum_iter
+        loss_scaler(loss, optimizer, parameters=model.parameters(),
+                    update_grad=(data_iter_step + 1) % accum_iter == 0)
+        if (data_iter_step + 1) % accum_iter == 0:
+            optimizer.zero_grad()
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(closs=c_loss_value)
+        metric_logger.update(mloss=m_loss_value)
+
+        lr = optimizer.param_groups[0]["lr"]
+        metric_logger.update(lr=lr)
+
+        loss_value_reduce = misc.all_reduce_mean(loss_value)
+        c_loss_value_reduce = misc.all_reduce_mean(c_loss_value)
+        m_loss_value_reduce = misc.all_reduce_mean(m_loss_value)
+        if log_writer is not None and (data_iter_step + 1) % accum_iter == 0:
+            """ We use epoch_1000x as the x-axis in tensorboard.
+            This calibrates different curves when batch size changes.
+            """
+            epoch_1000x = int((data_iter_step / len(data_loader) + epoch) * 1000)
+            log_writer.add_scalar('c_train_loss', c_loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('m_train_loss', m_loss_value_reduce, epoch_1000x)
+            log_writer.add_scalar('lr', lr, epoch_1000x)
+
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    return {k: meter.global_avg for k, meter in metric_logger.meters.items()}