Support dynamic resolution in evaluation

Author: czczup

internvl_chat/eval/caption/evaluate_caption.py

@@ -7,12 +7,13 @@
 from functools import partial
 
 import torch
-from internvl.train.dataset import build_transform
+from internvl.model.internvl_chat import InternVLChatModel
+from internvl.train.dataset import build_transform, dynamic_preprocess
 from PIL import Image
 from pycocoevalcap.eval import COCOEvalCap
 from pycocotools.coco import COCO
 from tqdm import tqdm
-from transformers import LlamaTokenizer
+from transformers import AutoTokenizer
 
 ds_collections = {
     'flickr30k': {
@@ -39,15 +40,20 @@
 
 class CaptionDataset(torch.utils.data.Dataset):
 
-    def __init__(self, name, root, annotation, prompt, input_size=224, pad2square=False):
+    def __init__(self, name, root, annotation, prompt, input_size=224, dynamic_image_size=False,
+                 use_thumbnail=False, max_num=6):
         if name == 'coco':
             self.images = json.load(open(annotation))
         else:
             self.images = json.load(open(annotation))['images']
         self.name = name
         self.prompt = prompt
         self.root = root
-        self.transform = build_transform(is_train=False, input_size=input_size, pad2square=pad2square)
+        self.input_size = input_size
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.max_num = max_num
+        self.transform = build_transform(is_train=False, input_size=input_size)
 
     def __len__(self):
         return len(self.images)
@@ -65,7 +71,14 @@ def __getitem__(self, idx):
                 image_path = os.path.join(self.root, self.images[idx]['image'])
 
         image = Image.open(image_path)
-        pixel_values = self.transform(image).unsqueeze(0)
+        if self.dynamic_image_size:
+            images = dynamic_preprocess(image, image_size=self.input_size,
+                                        use_thumbnail=self.use_thumbnail,
+                                        max_num=self.max_num)
+        else:
+            images = [image]
+        pixel_values = [self.transform(image) for image in images]
+        pixel_values = torch.stack(pixel_values)
 
         return {
             'image_id': image_id,
@@ -125,7 +138,9 @@ def evaluate_chat_model():
             annotation=annotation,
             prompt=prompt,
             input_size=image_size,
-            pad2square=pad2square
+            dynamic_image_size=args.dynamic,
+            use_thumbnail=use_thumbnail,
+            max_num=args.max_num
         )
         dataloader = torch.utils.data.DataLoader(
             dataset=dataset,
@@ -151,7 +166,7 @@ def evaluate_chat_model():
                 tokenizer=tokenizer,
                 pixel_values=pixel_values,
                 question=prompt,
-                generation_config=generation_config,
+                generation_config=generation_config
             )
             image_ids.extend(ids)
             captions.extend([pred])
@@ -217,6 +232,8 @@ def evaluate_chat_model():
     parser.add_argument('--temperature', type=float, default=0.0)
     parser.add_argument('--out-dir', type=str, default='results')
     parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--dynamic', action='store_true')
+    parser.add_argument('--max-num', type=int, default=6)
     args = parser.parse_args()
 
     if not os.path.exists(args.out_dir):
@@ -234,29 +251,22 @@ def evaluate_chat_model():
 
     torch.cuda.set_device(int(os.getenv('LOCAL_RANK', 0)))
 
-    tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
-
-    if 'qllama' in args.checkpoint.lower():
-        from internvl.model.internvl_chat_with_qllama import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.internvl.config.force_image_size or model.config.internvl_config.vision_config.image_size
-        pad2square = model.config.pad2square
-    else:
-        from internvl.model.internvl_chat import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.config.force_image_size or model.config.vision_config.image_size
-        pad2square = model.config.pad2square
+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False)
+    model = InternVLChatModel.from_pretrained(
+        args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
+    image_size = model.config.force_image_size or model.config.vision_config.image_size
+    use_thumbnail = model.config.use_thumbnail
 
     total_params = sum(p.numel() for p in model.parameters()) / 1e9
-    if total_params > 30:
+    if total_params > 20:
         args.num_beams = 1
         print(f'[test] total_params: {total_params}B, use num_beams: {args.num_beams}')
     else:
         print(f'[test] total_params: {total_params}B')
     print(f'[test] image_size: {image_size}')
-    print(f'[test] pad2square: {pad2square}')
     print(f'[test] template: {model.config.template}')
+    print(f'[test] dynamic_image_size: {args.dynamic}')
+    print(f'[test] use_thumbnail: {use_thumbnail}')
+    print(f'[test] max_num: {args.max_num}')
 
     evaluate_chat_model()

internvl_chat/eval/cmmmu/evaluate_cmmmu.py

@@ -4,10 +4,11 @@
 import random
 
 import torch
-from internvl.train.dataset import build_transform
+from internvl.model.internvl_chat import InternVLChatModel
+from internvl.train.dataset import build_transform, dynamic_preprocess
 from PIL import Image
 from tqdm import tqdm
-from transformers import LlamaTokenizer
+from transformers import AutoTokenizer
 
 ds_collections = {
     'art_and_design': {
@@ -51,15 +52,20 @@
 
 class VQADataset(torch.utils.data.Dataset):
 
-    def __init__(self, root, annotation, input_size=224, pad2square=False):
+    def __init__(self, root, annotation, input_size=224, dynamic_image_size=False,
+                 use_thumbnail=False, max_num=6):
         self.root = root
         self.items = []
         f = open(annotation)
         data = f.readlines()
         for data_line in data:
             data_line = json.loads(data_line)
             self.items.append(data_line)
-        self.transform = build_transform(is_train=False, input_size=input_size, pad2square=pad2square)
+        self.input_size = input_size
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.max_num = max_num
+        self.transform = build_transform(is_train=False, input_size=input_size)
 
     def __len__(self):
         return len(self.items)
@@ -69,7 +75,15 @@ def __getitem__(self, idx):
         image_path, question = item['image'], item['text']
         image_path = os.path.join(self.root, image_path)
         image = Image.open(image_path).convert('RGB')
-        pixel_values = self.transform(image).unsqueeze(0)
+        if self.dynamic_image_size:
+            images = dynamic_preprocess(image, image_size=self.input_size,
+                                        use_thumbnail=self.use_thumbnail,
+                                        max_num=self.max_num)
+        else:
+            images = [image]
+        pixel_values = [self.transform(image) for image in images]
+        pixel_values = torch.stack(pixel_values)
+
         return {
             'question': question,
             'pixel_values': pixel_values,
@@ -85,7 +99,9 @@ def evaluate_chat_model():
             root=ds_collections[ds_name]['root'],
             annotation=ds_collections[ds_name]['annotation'],
             input_size=image_size,
-            pad2square=pad2square
+            dynamic_image_size=args.dynamic,
+            use_thumbnail=use_thumbnail,
+            max_num=args.max_num
         )
 
         print(f'Evaluating {ds_name} ...')
@@ -109,9 +125,8 @@ def evaluate_chat_model():
                 tokenizer=tokenizer,
                 pixel_values=pixel_value,
                 question=question,
-                generation_config=generation_config,
+                generation_config=generation_config
             )
-            print(question, pred)
             question_id = item['question_id']
             text = item['text']
             output = {
@@ -137,6 +152,8 @@ def evaluate_chat_model():
     parser.add_argument('--temperature', type=float, default=0.0)
     parser.add_argument('--out-dir', type=str, default='results')
     parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--dynamic', action='store_true')
+    parser.add_argument('--max-num', type=int, default=6)
     args = parser.parse_args()
 
     if not os.path.exists(args.out_dir):
@@ -146,30 +163,23 @@ def evaluate_chat_model():
     print('datasets:', args.datasets)
     assert args.batch_size == 1, 'Only batch size 1 is supported'
 
-    tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
-
-    if 'qllama' in args.checkpoint.lower():
-        from internvl.model.internvl_chat_with_qllama import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.internvl.config.force_image_size or model.config.internvl_config.vision_config.image_size
-        pad2square = model.config.pad2square
-    else:
-        from internvl.model.internvl_chat import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.config.force_image_size or model.config.vision_config.image_size
-        pad2square = model.config.pad2square
+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False)
+    model = InternVLChatModel.from_pretrained(
+        args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
+    image_size = model.config.force_image_size or model.config.vision_config.image_size
+    use_thumbnail = model.config.use_thumbnail
 
     total_params = sum(p.numel() for p in model.parameters()) / 1e9
-    if total_params > 30:
+    if total_params > 20:
         args.num_beams = 1
         print(f'[test] total_params: {total_params}B, use num_beams: {args.num_beams}')
     else:
         print(f'[test] total_params: {total_params}B')
     print(f'[test] image_size: {image_size}')
-    print(f'[test] pad2square: {pad2square}')
     print(f'[test] template: {model.config.template}')
+    print(f'[test] dynamic_image_size: {args.dynamic}')
+    print(f'[test] use_thumbnail: {use_thumbnail}')
+    print(f'[test] max_num: {args.max_num}')
 
     model_id = '_'.join(args.checkpoint.split('/')[-2:])
     evaluate_chat_model()

internvl_chat/eval/llava_bench/evaluate_llava_bench.py

@@ -4,10 +4,11 @@
 import random
 
 import torch
-from internvl.train.dataset import build_transform
+from internvl.model.internvl_chat import InternVLChatModel
+from internvl.train.dataset import build_transform, dynamic_preprocess
 from PIL import Image
 from tqdm import tqdm
-from transformers import LlamaTokenizer
+from transformers import AutoTokenizer
 
 ds_collections = {
     'llava_bench': {
@@ -21,11 +22,16 @@
 
 class VQADataset(torch.utils.data.Dataset):
 
-    def __init__(self, root, data, prompt, input_size=224, pad2square=False):
+    def __init__(self, root, data, prompt, input_size=224, dynamic_image_size=False,
+                 use_thumbnail=False, max_num=6):
         self.root = root
         self.data = open(data).readlines()
         self.prompt = prompt
-        self.transform = build_transform(is_train=False, input_size=input_size, pad2square=pad2square)
+        self.input_size = input_size
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.max_num = max_num
+        self.transform = build_transform(is_train=False, input_size=input_size)
 
     def __len__(self):
         return len(self.data)
@@ -37,7 +43,14 @@ def __getitem__(self, idx):
 
         image = os.path.join(self.root, image)
         image = Image.open(image).convert('RGB')
-        pixel_values = self.transform(image).unsqueeze(0)
+        if self.dynamic_image_size:
+            images = dynamic_preprocess(image, image_size=self.input_size,
+                                        use_thumbnail=self.use_thumbnail,
+                                        max_num=self.max_num)
+        else:
+            images = [image]
+        pixel_values = [self.transform(image) for image in images]
+        pixel_values = torch.stack(pixel_values)
         question = question + self.prompt
         return question_id, question, pixel_values, annotation
 
@@ -51,7 +64,9 @@ def evaluate_chat_model():
             data=ds_collections[ds_name]['question'],
             prompt=' Please give a detailed answer.',
             input_size=image_size,
-            pad2square=pad2square
+            dynamic_image_size=args.dynamic,
+            use_thumbnail=use_thumbnail,
+            max_num=args.max_num
         )
 
         outputs = []
@@ -61,16 +76,14 @@ def evaluate_chat_model():
                 num_beams=args.num_beams,
                 max_new_tokens=ds_collections[ds_name]['max_new_tokens'],
                 min_new_tokens=ds_collections[ds_name]['min_new_tokens'],
-                length_penalty=1,
-                # repetition_penalty=1.5,
                 do_sample=True if args.temperature > 0 else False,
                 temperature=args.temperature,
             )
             pred = model.chat(
                 tokenizer=tokenizer,
                 pixel_values=pixel_values,
                 question=question,
-                generation_config=generation_config,
+                generation_config=generation_config
             )
             outputs.append({
                 'question_id': question_id,
@@ -100,6 +113,8 @@ def evaluate_chat_model():
     parser.add_argument('--temperature', type=float, default=0.0)
     parser.add_argument('--out-dir', type=str, default='results')
     parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--dynamic', action='store_true')
+    parser.add_argument('--max-num', type=int, default=6)
     args = parser.parse_args()
 
     if not os.path.exists(args.out_dir):
@@ -109,30 +124,23 @@ def evaluate_chat_model():
     print('datasets:', args.datasets)
     assert args.batch_size == 1, 'Only batch size 1 is supported'
 
-    tokenizer = LlamaTokenizer.from_pretrained(args.checkpoint)
-
-    if 'qllama' in args.checkpoint.lower():
-        from internvl.model.internvl_chat_with_qllama import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.internvl.config.force_image_size or model.config.internvl_config.vision_config.image_size
-        pad2square = model.config.pad2square
-    else:
-        from internvl.model.internvl_chat import InternVLChatModel
-        model = InternVLChatModel.from_pretrained(
-            args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
-        image_size = model.config.force_image_size or model.config.vision_config.image_size
-        pad2square = model.config.pad2square
+    tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, trust_remote_code=True, use_fast=False)
+    model = InternVLChatModel.from_pretrained(
+        args.checkpoint, low_cpu_mem_usage=True, torch_dtype=torch.bfloat16).cuda().eval()
+    image_size = model.config.force_image_size or model.config.vision_config.image_size
+    use_thumbnail = model.config.use_thumbnail
 
     total_params = sum(p.numel() for p in model.parameters()) / 1e9
-    if total_params > 30:
+    if total_params > 20:
         args.num_beams = 1
         print(f'[test] total_params: {total_params}B, use num_beams: {args.num_beams}')
     else:
         print(f'[test] total_params: {total_params}B')
     print(f'[test] image_size: {image_size}')
-    print(f'[test] pad2square: {pad2square}')
     print(f'[test] template: {model.config.template}')
+    print(f'[test] dynamic_image_size: {args.dynamic}')
+    print(f'[test] use_thumbnail: {use_thumbnail}')
+    print(f'[test] max_num: {args.max_num}')
 
     model_id = '_'.join(args.checkpoint.split('/')[-2:])
     evaluate_chat_model()