From 0c79b57fcccd264e004ee73d5ebdbec970dd16d5 Mon Sep 17 00:00:00 2001 From: "yichang.zyc" Date: Wed, 28 Dec 2022 12:17:36 +0800 Subject: [PATCH] support batch infer Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11170755 --- modelscope/metrics/accuracy_metric.py | 3 + .../models/multi_modal/ofa_for_all_tasks.py | 52 +++++++++------ modelscope/pipelines/base.py | 19 +++--- .../cv/image_classification_pipeline.py | 12 ++++ .../pipelines/multi_modal/asr_pipeline.py | 8 ++- .../multi_modal/image_captioning_pipeline.py | 16 +---- .../multi_modal/ocr_recognition_pipeline.py | 8 ++- .../multi_modal/visual_entailment_pipeline.py | 14 ++++ .../multi_modal/visual_grounding_pipeline.py | 14 ++++ .../visual_question_answering_pipeline.py | 8 ++- .../pipelines/nlp/summarization_pipeline.py | 14 ++++ .../nlp/text_classification_pipeline.py | 11 +++- modelscope/pipelines/util.py | 24 +++++++ .../preprocessors/ofa/image_classification.py | 3 +- .../preprocessors/ofa/text_classification.py | 3 + .../preprocessors/ofa/visual_entailment.py | 5 +- tests/pipelines/test_ofa_tasks.py | 64 ++++++++++++++++--- tests/trainers/test_ofa_trainer.py | 2 +- 18 files changed, 224 insertions(+), 56 deletions(-) diff --git a/modelscope/metrics/accuracy_metric.py b/modelscope/metrics/accuracy_metric.py index fe040177..f6c3176b 100644 --- a/modelscope/metrics/accuracy_metric.py +++ b/modelscope/metrics/accuracy_metric.py @@ -40,6 +40,9 @@ class AccuracyMetric(Metric): self.labels.append(truth) for result in eval_results: if isinstance(truth, str): + if isinstance(result, list): + result = result[0] + assert isinstance(result, str), 'both truth and pred are str' self.preds.append(remove_space_between_chinese_chars(result)) else: self.preds.append(result) diff --git a/modelscope/models/multi_modal/ofa_for_all_tasks.py b/modelscope/models/multi_modal/ofa_for_all_tasks.py index 3a35be58..7f459b0c 100644 --- a/modelscope/models/multi_modal/ofa_for_all_tasks.py +++ b/modelscope/models/multi_modal/ofa_for_all_tasks.py @@ -105,6 +105,8 @@ class OfaForAllTasks(TorchModel): } if hasattr(self.cfg.model, 'beam_search'): sg_args.update(self.cfg.model.beam_search) + self.num_return_sequences = self.cfg.model.get('num_return_sequences', + 1) if len(self.ans2label_dict) > 0: self.constraint_trie = Trie(self.tokenizer.eos_token_id) self.val_ans_l = [] @@ -140,15 +142,14 @@ class OfaForAllTasks(TorchModel): return self.inference(input) def inference(self, input: Dict[str, Any]) -> Dict[str, Any]: + assert self.generator.beam_size >= self.num_return_sequences, \ + 'beam search can only return beam size sentences' + if self.ans2label_dict and self.gen_type == 'generation': + assert self.generator.beam_size <= len(self.ans2label_dict), \ + 'beam search will not work properly.' ret = self.task_inference_mapping[self.cfg.task](input) if 'samples' in input: ret['samples'] = input['samples'] - for key in [ - OutputKeys.CAPTION, OutputKeys.TEXT, OutputKeys.BOXES, - OutputKeys.LABELS, OutputKeys.SCORES - ]: - if key not in ret: - ret[key] = None return ret def postprocess(self, input: Dict[str, Any], **kwargs) -> Dict[str, Any]: @@ -157,7 +158,8 @@ class OfaForAllTasks(TorchModel): result_l = list() for cap in caption: if self.language == 'en': - result_l.append(cap.translate(self.transtab).strip()) + result_l.append( + [c.translate(self.transtab).strip() for c in cap]) else: result_l.append(cap) input[OutputKeys.CAPTION] = result_l @@ -166,8 +168,18 @@ class OfaForAllTasks(TorchModel): ] and self.cfg.task != Tasks.visual_grounding: ret_l = list() for text in input[OFA_TASK_KEY_MAPPING[self.cfg.task]]: - ret_l.append(self.detokenizer(text)) + ret_l.append([self.detokenizer(t) for t in text]) input[OFA_TASK_KEY_MAPPING[self.cfg.task]] = ret_l + for key in [ + OutputKeys.CAPTION, OutputKeys.TEXT, OutputKeys.BOXES, + OutputKeys.LABELS, OutputKeys.SCORES + ]: + if key not in input: + input[key] = None + else: + if (len(input[key]) == 1 and isinstance(input[key], list)) \ + and self.cfg.task != Tasks.visual_grounding: + input[key] = input[key][0] return input def _text_gen_inference(self, input): @@ -175,23 +187,25 @@ class OfaForAllTasks(TorchModel): input, prefix_tokens=input.get( 'prefix_tokens', None)) - gen_l = list() + results = list() for idx, gen_out in enumerate(gen_outputs): - if len(gen_out) > 0: - decode_tokens = gen_out[0]['tokens'] + gen_token_l = [] + for beam_gen_out in gen_out[:self.num_return_sequences]: + decode_tokens = beam_gen_out['tokens'] if 'prefix_tokens' in input: prefix_len = input['prefix_tokens'][idx].ne( self.pad_item.to(self.model.device)).sum() decode_tokens = decode_tokens[prefix_len:] - gen_l.append(decode_tokens) - else: - gen_l.append('') - result = self.tokenizer.batch_decode(gen_l, skip_special_tokens=True) - result = [item.strip() for item in result] + gen_token_l.append(decode_tokens) + result = self.tokenizer.batch_decode( + gen_token_l, skip_special_tokens=True) + result = [item.strip() for item in result] + result.extend([''] * (self.num_return_sequences - len(result))) + results.append(result) # text generation tasks have no score - ret = {OFA_TASK_KEY_MAPPING[self.cfg.task]: result} - if self.cfg.task.endswith('classification'): - ret[OutputKeys.SCORES] = [1.0] * len(result) + ret = {OFA_TASK_KEY_MAPPING[self.cfg.task]: results} + if self.ans2label_dict: + ret[OutputKeys.SCORES] = [[1.0]] * len(results) return ret def _visual_grounding_inference(self, input): diff --git a/modelscope/pipelines/base.py b/modelscope/pipelines/base.py index c49fe7dc..4863a290 100644 --- a/modelscope/pipelines/base.py +++ b/modelscope/pipelines/base.py @@ -273,14 +273,17 @@ class Pipeline(ABC): **forward_params) else: batched_out = self.forward(batched_input, **forward_params) - for batch_idx in range(real_batch_size): - out = {} - for k, element in batched_out.items(): - if element is not None: - out[k] = element[batch_idx] - out = self.postprocess(out, **postprocess_params) - self._check_output(out) - output_list.append(out) + if real_batch_size == 1: + output_list.append(batched_out) + else: + for batch_idx in range(real_batch_size): + out = {} + for k, element in batched_out.items(): + if element is not None: + out[k] = element[batch_idx] + out = self.postprocess(out, **postprocess_params) + self._check_output(out) + output_list.append(out) return output_list diff --git a/modelscope/pipelines/cv/image_classification_pipeline.py b/modelscope/pipelines/cv/image_classification_pipeline.py index b9d7376b..52b27141 100644 --- a/modelscope/pipelines/cv/image_classification_pipeline.py +++ b/modelscope/pipelines/cv/image_classification_pipeline.py @@ -11,6 +11,7 @@ from modelscope.models.multi_modal import OfaForAllTasks from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Input, Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import OfaPreprocessor, Preprocessor, load_image from modelscope.utils.constant import Tasks from modelscope.utils.device import get_device @@ -35,6 +36,17 @@ class ImageClassificationPipeline(Pipeline): if preprocessor is None and isinstance(self.model, OfaForAllTasks): self.preprocessor = OfaPreprocessor(model_dir=self.model.model_dir) + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(ImageClassificationPipeline, self)._batch(data) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + with torch.no_grad(): + return super().forward(inputs, **forward_params) + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: return inputs diff --git a/modelscope/pipelines/multi_modal/asr_pipeline.py b/modelscope/pipelines/multi_modal/asr_pipeline.py index 3cb7439c..590a2496 100644 --- a/modelscope/pipelines/multi_modal/asr_pipeline.py +++ b/modelscope/pipelines/multi_modal/asr_pipeline.py @@ -5,9 +5,9 @@ import torch from modelscope.metainfo import Pipelines from modelscope.models.multi_modal import MPlugForAllTasks, OfaForAllTasks -from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import (MPlugPreprocessor, OfaPreprocessor, Preprocessor) from modelscope.utils.constant import Tasks @@ -45,6 +45,12 @@ class AutomaticSpeechRecognitionPipeline(Pipeline): preprocessor = MPlugPreprocessor(pipe_model.model_dir) super().__init__(model=pipe_model, preprocessor=preprocessor, **kwargs) + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(AutomaticSpeechRecognitionPipeline, self)._batch(data) + def forward(self, inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: with torch.no_grad(): diff --git a/modelscope/pipelines/multi_modal/image_captioning_pipeline.py b/modelscope/pipelines/multi_modal/image_captioning_pipeline.py index e1d5c769..0a16c58f 100644 --- a/modelscope/pipelines/multi_modal/image_captioning_pipeline.py +++ b/modelscope/pipelines/multi_modal/image_captioning_pipeline.py @@ -5,9 +5,9 @@ import torch from modelscope.metainfo import Pipelines from modelscope.models.multi_modal import MPlugForAllTasks, OfaForAllTasks -from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import (MPlugPreprocessor, OfaPreprocessor, Preprocessor) from modelscope.utils.constant import Tasks @@ -39,17 +39,7 @@ class ImageCaptioningPipeline(Pipeline): def _batch(self, data): if isinstance(self.model, OfaForAllTasks): - # collate batch data due to the nested data structure - if isinstance(data, list): - batch_data = {} - batch_data['nsentences'] = len(data) - batch_data['samples'] = [d['samples'][0] for d in data] - batch_data['net_input'] = {} - for k in data[0]['net_input'].keys(): - batch_data['net_input'][k] = torch.cat( - [d['net_input'][k] for d in data]) - - return batch_data + return batch_process(self.model, data) elif isinstance(self.model, MPlugForAllTasks): from transformers.tokenization_utils_base import BatchEncoding batch_data = dict(train=data[0]['train']) @@ -60,7 +50,7 @@ class ImageCaptioningPipeline(Pipeline): batch_data['question'] = BatchEncoding(question) return batch_data else: - return super()._collate_batch(data) + return super(ImageCaptioningPipeline, self)._batch(data) def forward(self, inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: diff --git a/modelscope/pipelines/multi_modal/ocr_recognition_pipeline.py b/modelscope/pipelines/multi_modal/ocr_recognition_pipeline.py index 3c4a3c3c..337742a7 100644 --- a/modelscope/pipelines/multi_modal/ocr_recognition_pipeline.py +++ b/modelscope/pipelines/multi_modal/ocr_recognition_pipeline.py @@ -5,9 +5,9 @@ import torch from modelscope.metainfo import Pipelines from modelscope.models.multi_modal import OfaForAllTasks -from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import OfaPreprocessor, Preprocessor from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger @@ -34,6 +34,12 @@ class OcrRecognitionPipeline(Pipeline): if isinstance(self.model, OfaForAllTasks): self.preprocessor = OfaPreprocessor(self.model.model_dir) + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(OcrRecognitionPipeline, self)._batch(data) + def forward(self, inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: with torch.no_grad(): diff --git a/modelscope/pipelines/multi_modal/visual_entailment_pipeline.py b/modelscope/pipelines/multi_modal/visual_entailment_pipeline.py index 67661b39..703b8d1b 100644 --- a/modelscope/pipelines/multi_modal/visual_entailment_pipeline.py +++ b/modelscope/pipelines/multi_modal/visual_entailment_pipeline.py @@ -1,10 +1,13 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from typing import Any, Dict, Optional, Union +import torch + from modelscope.metainfo import Pipelines from modelscope.models.multi_modal import OfaForAllTasks from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import OfaPreprocessor, Preprocessor from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger @@ -30,5 +33,16 @@ class VisualEntailmentPipeline(Pipeline): if preprocessor is None and isinstance(self.model, OfaForAllTasks): self.preprocessor = OfaPreprocessor(model_dir=self.model.model_dir) + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(VisualEntailmentPipeline, self)._batch(data) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + with torch.no_grad(): + return super().forward(inputs, **forward_params) + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: return inputs diff --git a/modelscope/pipelines/multi_modal/visual_grounding_pipeline.py b/modelscope/pipelines/multi_modal/visual_grounding_pipeline.py index f8a79d55..36b6754b 100644 --- a/modelscope/pipelines/multi_modal/visual_grounding_pipeline.py +++ b/modelscope/pipelines/multi_modal/visual_grounding_pipeline.py @@ -1,10 +1,13 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from typing import Any, Dict, Optional, Union +import torch + from modelscope.metainfo import Pipelines from modelscope.models.multi_modal import OfaForAllTasks from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import OfaPreprocessor, Preprocessor from modelscope.utils.constant import Tasks from modelscope.utils.logger import get_logger @@ -30,5 +33,16 @@ class VisualGroundingPipeline(Pipeline): if preprocessor is None and isinstance(self.model, OfaForAllTasks): self.preprocessor = OfaPreprocessor(model_dir=self.model.model_dir) + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(VisualGroundingPipeline, self)._batch(data) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + with torch.no_grad(): + return super().forward(inputs, **forward_params) + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: return inputs diff --git a/modelscope/pipelines/multi_modal/visual_question_answering_pipeline.py b/modelscope/pipelines/multi_modal/visual_question_answering_pipeline.py index a30cf1c5..09b6e1ba 100644 --- a/modelscope/pipelines/multi_modal/visual_question_answering_pipeline.py +++ b/modelscope/pipelines/multi_modal/visual_question_answering_pipeline.py @@ -6,9 +6,9 @@ import torch from modelscope.metainfo import Pipelines from modelscope.models import Model from modelscope.models.multi_modal import MPlugForAllTasks, OfaForAllTasks -from modelscope.outputs import OutputKeys from modelscope.pipelines.base import Pipeline, Tensor from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import (MPlugPreprocessor, OfaPreprocessor, Preprocessor) from modelscope.utils.constant import Tasks @@ -39,6 +39,12 @@ class VisualQuestionAnsweringPipeline(Pipeline): self.preprocessor = MPlugPreprocessor(self.model.model_dir) self.model.eval() + def _batch(self, data): + if isinstance(self.model, OfaForAllTasks): + return batch_process(self.model, data) + else: + return super(VisualQuestionAnsweringPipeline, self)._batch(data) + def forward(self, inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: with torch.no_grad(): diff --git a/modelscope/pipelines/nlp/summarization_pipeline.py b/modelscope/pipelines/nlp/summarization_pipeline.py index 7c8355f9..5c5e5305 100644 --- a/modelscope/pipelines/nlp/summarization_pipeline.py +++ b/modelscope/pipelines/nlp/summarization_pipeline.py @@ -1,9 +1,12 @@ # Copyright (c) Alibaba, Inc. and its affiliates. from typing import Any, Dict, Optional, Union +import torch + from modelscope.metainfo import Pipelines, Preprocessors from modelscope.pipelines.base import Model, Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import Preprocessor from modelscope.utils.constant import Fields, Tasks from modelscope.utils.logger import get_logger @@ -48,5 +51,16 @@ class SummarizationPipeline(Pipeline): self.preprocessor = Preprocessor.from_pretrained( self.model.model_dir, **kwargs) + def _batch(self, data): + if self.model.__class__.__name__ == 'OfaForAllTasks': + return batch_process(self.model, data) + else: + return super(SummarizationPipeline, self)._batch(data) + + def forward(self, inputs: Dict[str, Any], + **forward_params) -> Dict[str, Any]: + with torch.no_grad(): + return super().forward(inputs, **forward_params) + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: return inputs diff --git a/modelscope/pipelines/nlp/text_classification_pipeline.py b/modelscope/pipelines/nlp/text_classification_pipeline.py index 06f95d7e..9a3b6901 100644 --- a/modelscope/pipelines/nlp/text_classification_pipeline.py +++ b/modelscope/pipelines/nlp/text_classification_pipeline.py @@ -2,12 +2,14 @@ from typing import Any, Dict, Union import numpy as np +import torch from modelscope.metainfo import Pipelines, Preprocessors from modelscope.models.base import Model from modelscope.outputs import OutputKeys, TextClassificationModelOutput from modelscope.pipelines.base import Pipeline from modelscope.pipelines.builder import PIPELINES +from modelscope.pipelines.util import batch_process from modelscope.preprocessors import Preprocessor from modelscope.utils.constant import Fields, Tasks from modelscope.utils.logger import get_logger @@ -83,10 +85,17 @@ class TextClassificationPipeline(Pipeline): if hasattr(self.preprocessor, 'id2label'): self.id2label = self.preprocessor.id2label + def _batch(self, data): + if self.model.__class__.__name__ == 'OfaForAllTasks': + return batch_process(self.model, data) + else: + return super(TextClassificationPipeline, self)._batch(data) + def forward(self, inputs: Dict[str, Any], **forward_params) -> Dict[str, Any]: if self.model.__class__.__name__ == 'OfaForAllTasks': - return super().forward(inputs, **forward_params) + with torch.no_grad(): + return super().forward(inputs, **forward_params) return self.model(**inputs, **forward_params) def postprocess(self, diff --git a/modelscope/pipelines/util.py b/modelscope/pipelines/util.py index 99a11317..fbbf4084 100644 --- a/modelscope/pipelines/util.py +++ b/modelscope/pipelines/util.py @@ -2,6 +2,8 @@ import os.path as osp from typing import List, Optional, Union +import torch + from modelscope.hub.api import HubApi from modelscope.hub.file_download import model_file_download from modelscope.utils.config import Config @@ -81,3 +83,25 @@ def is_model(path: Union[str, List]): ) return all_true + + +def batch_process(model, data): + if model.__class__.__name__ == 'OfaForAllTasks': + # collate batch data due to the nested data structure + assert isinstance(data, list) + batch_data = { + 'nsentences': len(data), + 'samples': [d['samples'][0] for d in data], + 'net_input': {} + } + for k in data[0]['net_input'].keys(): + batch_data['net_input'][k] = torch.cat( + [d['net_input'][k] for d in data]) + if 'w_resize_ratios' in data[0]: + batch_data['w_resize_ratios'] = torch.cat( + [d['w_resize_ratios'] for d in data]) + if 'h_resize_ratios' in data[0]: + batch_data['h_resize_ratios'] = torch.cat( + [d['h_resize_ratios'] for d in data]) + + return batch_data diff --git a/modelscope/preprocessors/ofa/image_classification.py b/modelscope/preprocessors/ofa/image_classification.py index 038a9e15..f3f44647 100644 --- a/modelscope/preprocessors/ofa/image_classification.py +++ b/modelscope/preprocessors/ofa/image_classification.py @@ -112,7 +112,8 @@ class OfaImageClassificationPreprocessor(OfaBasePreprocessor): sample = { 'source': inputs, 'patch_image': patch_image, - 'patch_mask': torch.tensor([True]) + 'patch_mask': torch.tensor([True]), + 'decoder_prompt': self.bos_item, } if 'text' in self.column_map and self.column_map['text'] in data: sample['label'] = data[self.column_map['text']] diff --git a/modelscope/preprocessors/ofa/text_classification.py b/modelscope/preprocessors/ofa/text_classification.py index 24c4f67e..890d5040 100644 --- a/modelscope/preprocessors/ofa/text_classification.py +++ b/modelscope/preprocessors/ofa/text_classification.py @@ -68,13 +68,16 @@ class OfaTextClassificationPreprocessor(OfaBasePreprocessor): instruction_itm = self._build_instruction(data) if self.prompt_type == 'none': prefix_token = [] + decoder_prompt = self.bos_item elif self.prompt_type == 'prev_output': prefix_token = instruction_itm[:-1] # remove eos + decoder_prompt = instruction_itm[:-1] else: raise NotImplementedError sample = { 'source': instruction_itm, 'prefix_token': prefix_token, + 'decoder_prompt': decoder_prompt, } if 'label' in data: sample['label'] = self.label2ans[data['label']] diff --git a/modelscope/preprocessors/ofa/visual_entailment.py b/modelscope/preprocessors/ofa/visual_entailment.py index fff5bbd3..468c9dc7 100644 --- a/modelscope/preprocessors/ofa/visual_entailment.py +++ b/modelscope/preprocessors/ofa/visual_entailment.py @@ -101,10 +101,10 @@ class OfaVisualEntailmentPreprocessor(OfaBasePreprocessor): text = prompt.format(caption, hypothesis) inputs = self.tokenize_text(text) if self.prompt_type == 'none': + prefix_token = [] decoder_prompt = self.bos_item - elif self.prompt_type == 'src': - decoder_prompt = inputs elif self.prompt_type == 'prev_output': + prefix_token = inputs[:-1] # remove eos decoder_prompt = inputs[:-1] else: raise NotImplementedError @@ -112,6 +112,7 @@ class OfaVisualEntailmentPreprocessor(OfaBasePreprocessor): 'source': inputs, 'patch_image': patch_image, 'patch_mask': torch.tensor([True]), + 'prefix_token': prefix_token, 'decoder_prompt': decoder_prompt, } if 'relation' in self.column_map and self.column_map[ diff --git a/tests/pipelines/test_ofa_tasks.py b/tests/pipelines/test_ofa_tasks.py index 6dec2c57..0690424c 100644 --- a/tests/pipelines/test_ofa_tasks.py +++ b/tests/pipelines/test_ofa_tasks.py @@ -45,15 +45,16 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): result = img_captioning('data/test/images/image_captioning.png') print(result[OutputKeys.CAPTION]) - @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') - def test_run_with_image_captioning_batch(self): - img_captioning = pipeline( - Tasks.image_captioning, - model='damo/ofa_image-caption_coco_large_en') + img_captioning.model.num_return_sequences = 2 + result = img_captioning('data/test/images/image_captioning.png') + print(result[OutputKeys.CAPTION]) + + # test batch infer + img_captioning.model.num_return_sequences = 1 results = img_captioning( [{ 'image': 'data/test/images/image_captioning.png' - } for _ in range(6)], + } for _ in range(3)], batch_size=2) for r in results: print(r[OutputKeys.CAPTION]) @@ -65,6 +66,12 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): model='damo/ofa_ocr-recognition_scene_base_zh') result = ocr_recognize('data/test/images/image_ocr_recognition.jpg') print(result[OutputKeys.TEXT]) + # test batch infer + results = ocr_recognize( + ['data/test/images/image_ocr_recognition.jpg' for _ in range(3)], + batch_size=2) + for r in results: + print(r[OutputKeys.TEXT]) @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_image_classification_with_model(self): @@ -84,6 +91,12 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): result = ofa_pipe(image) print(result) + # test batch infer + image = ['data/test/images/image_classification.png' for _ in range(3)] + results = ofa_pipe(image, batch_size=2) + for r in results: + print(r[OutputKeys.LABELS], r[OutputKeys.SCORES]) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_summarization_with_model(self): model = Model.from_pretrained( @@ -104,12 +117,23 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): model='damo/ofa_summarization_gigaword_large_en') text = 'five-time world champion michelle kwan withdrew' + \ 'from the #### us figure skating championships on wednesday ,' + \ - ' but will petition us skating officials for the chance to ' +\ + ' but will petition us skating officials for the chance to ' + \ 'compete at the #### turin olympics .' input = {'text': text} result = ofa_pipe(input) print(result) + # test for return multiple sequences + ofa_pipe.model.num_return_sequences = 2 + result = ofa_pipe(input) + print(result) + # test batch infer + ofa_pipe.model.num_return_sequences = 1 + input = [{'text': text} for _ in range(3)] + results = ofa_pipe(input, batch_size=2) + for r in results: + print(r[OutputKeys.TEXT]) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_text_classification_with_model(self): model = Model.from_pretrained( @@ -130,6 +154,11 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): text2 = 'A member of my team will execute your orders with immense precision.' result = ofa_pipe((text, text2)) print(result) + # test batch infer + inputs = [(text, text2) for _ in range(3)] + results = ofa_pipe(inputs, batch_size=2) + for r in results: + print(r[OutputKeys.LABELS], r[OutputKeys.SCORES]) @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_visual_entailment_with_model(self): @@ -152,8 +181,13 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): input = {'image': image, 'text': text} result = ofa_pipe(input) print(result) + # test batch infer + input = [{'image': image, 'text': text} for _ in range(3)] + results = ofa_pipe(input, batch_size=2) + for r in results: + print(r[OutputKeys.LABELS], r[OutputKeys.SCORES]) - @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_with_visual_grounding_with_model(self): model = Model.from_pretrained( 'damo/ofa_visual-grounding_refcoco_large_en') @@ -182,6 +216,9 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): image_name = image.split('/')[-2] self.save_img(image, result[OutputKeys.BOXES][0], osp.join('large_en_name_' + image_name + '.png')) + # test batch infer + result = ofa_pipe([input for _ in range(3)], batch_size=2) + print(result) @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_with_visual_grounding_zh_with_name(self): @@ -217,6 +254,10 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): result = ofa_pipe(input) print(result) + # test batch infer + result = ofa_pipe([input for _ in range(3)], batch_size=2) + print(result) + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') def test_run_with_image_captioning_distilled_with_model(self): model = Model.from_pretrained( @@ -230,6 +271,9 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): result = img_captioning(image) print(result[OutputKeys.CAPTION]) + # test batch infer + print(img_captioning([image for _ in range(3)], batch_size=2)) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_with_visual_entailment_distilled_model_with_name(self): ofa_pipe = pipeline( @@ -280,6 +324,10 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck): example = {'wav': 'data/test/audios/asr_example_ofa.wav'} result = ofa_pipe(example) print(result[OutputKeys.TEXT]) + # test batch infer + result = ofa_pipe([example for _ in range(3)], batch_size=2) + for r in result: + print(r[OutputKeys.TEXT]) @unittest.skip('demo compatibility test is only enabled on a needed-basis') def test_demo_compatibility(self): diff --git a/tests/trainers/test_ofa_trainer.py b/tests/trainers/test_ofa_trainer.py index ab2b8cc6..f4ca7bcb 100644 --- a/tests/trainers/test_ofa_trainer.py +++ b/tests/trainers/test_ofa_trainer.py @@ -37,7 +37,7 @@ class TestOfaTrainer(unittest.TestCase): 'train': {'work_dir': 'work/ckpts/recognition', # 'launcher': 'pytorch', 'max_epochs': 1, - 'use_fp16': True, + 'use_fp16': False, 'dataloader': {'batch_size_per_gpu': 4, 'workers_per_gpu': 0}, 'lr_scheduler': {'name': 'polynomial_decay', 'warmup_proportion': 0.01,