diff --git a/__init__.py b/__init__.py index 8a0ea78..70f6cad 100644 --- a/__init__.py +++ b/__init__.py @@ -15,5 +15,5 @@ from .clip import Clip -def clip(model_name: str, modality: str): - return Clip(model_name, modality) +def clip(model_name: str, modality: str, device:str = None, checkpoint_path:str = None): + return Clip(model_name, modality, device, checkpoint_path) diff --git a/clip.py b/clip.py index 1539434..2041191 100644 --- a/clip.py +++ b/clip.py @@ -11,10 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - import sys from pathlib import Path import torch +from torch import nn from torchvision import transforms from towhee.types.image_utils import to_pil @@ -22,29 +22,55 @@ from towhee.operator.base import NNOperator, OperatorFlag from towhee.types.arg import arg, to_image_color from towhee import register from transformers import CLIPTokenizer, CLIPTextModel ,CLIPModel,CLIPProcessor -from train_clip_with_hf_trainer import train_with_hf_trainer +#@accelerate +class CLIPModelVision(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, pixel_values): + image_embeds = self.model.get_image_features(pixel_values) + return image_embeds + +#@accelerate +class CLIPModelText(nn.Module): + def __init__(self, model): + super().__init__() + self.model = model + + def forward(self, input_ids, attention_mask): + text_embeds = self.model.get_text_features(input_ids, attention_mask) + return text_embeds + @register(output_schema=['vec']) class Clip(NNOperator): """ CLIP multi-modal embedding operator """ - def __init__(self, model_name: str, modality: str): + def __init__(self, model_name: str, modality: str, device, checkpoint_path): + self.model_name = model_name self.modality = modality self.device = "cuda" if torch.cuda.is_available() else "cpu" cfg = self._configs()[model_name] - self.model = CLIPModel.from_pretrained(cfg['name']) - self.tokenizer = CLIPTokenizer.from_pretrained(cfg['name']) - self.processor = CLIPProcessor.from_pretrained(cfg['name']) + clip_model = CLIPModel.from_pretrained(cfg) + if self.modality == 'image': + self.model = CLIPModelVision(clip_model) + elif self.modality == 'text': + self.model = CLIPModelText(clip_model) + else: + raise ValueError("modality[{}] not implemented.".format(self.modality)) + self.tokenizer = CLIPTokenizer.from_pretrained(cfg) + self.processor = CLIPProcessor.from_pretrained(cfg) def inference_single_data(self, data): if self.modality == 'image': vec = self._inference_from_image(data) - elif self.modality == 'text': + elif self.modality == 'text': vec = self._inference_from_text(data) else: - raise ValueError("modality[{}] not implemented.".format(self._modality)) + raise ValueError("modality[{}] not implemented.".format(self.modality)) return vec.detach().cpu().numpy().flatten() def __call__(self, data): @@ -63,29 +89,122 @@ class Clip(NNOperator): def _inference_from_text(self, text): tokens = self.tokenizer([text], padding=True, return_tensors="pt") - text_features = self.model.get_text_features(**tokens) + text_features = self.model(tokens['input_ids'],tokens['attention_mask']) return text_features @arg(1, to_image_color('RGB')) def _inference_from_image(self, img): img = to_pil(img) - inputs = processor(images=img, return_tensors="pt") - image_features = self.model.get_image_features(**inputs) + inputs = self.processor(images=img, return_tensors="pt") + image_features = self.model(inputs['pixel_values']) return image_features def train(self, **kwargs): + import sys + import pathlib + path = str(pathlib.Path(__file__).parent) + print(path) + sys.path.append(path) + from train_clip_with_hf_trainer import train_with_hf_trainer data_args = kwargs.pop('data_args', None) training_args = kwargs.pop('training_args', None) train_with_hf_trainer(self.model, self.tokenizer, data_args, training_args) def _configs(self): config = {} - config['clip_vit_base_32'] = {} - config['clip_vit_base_32']['name'] = 'openai/clip-vit-base-patch16' - config['clip_vit_base_16'] = {} - config['clip_vit_base_16']['name'] = 'openai/clip-vit-base-patch32' - config['clip_vit_large_14'] = {} + config['clip_vit_base_32'] = 'openai/clip-vit-base-patch16' + config['clip_vit_base_16'] = 'openai/clip-vit-base-patch32' config['clip_vit_large_14'] = 'openai/clip-vit-large-patch14' - config['clip_vit_large_14_336'] = {} - config['clip_vit_large_14_336']['name'] ='openai/clip-vit-large-patch14-336' + config['clip_vit_large_14_336'] ='openai/clip-vit-large-patch14-336' return config + + @property + def supported_formats(self): + onnxes = self.supported_model_names(format='onnx') + if self.model_name in onnxes: + return ['onnx'] + else: + return ['pytorch'] + + @staticmethod + def supported_model_names(format: str = None): + if format == 'pytorch' or format == 'torchscript' or format == 'onnx': + model_list = [ + 'openai/clip-vit-base-patch16', + 'openai/clip-vit-base-patch32', + 'openai/clip-vit-large-patch14', + 'openai/clip-vit-large-patch14-336' + ] + else: + log.error(f'Invalid format "{format}". Currently supported formats: "pytorch", "torchscript".') + return model_list + + @property + def _model(self): + return self.model + + def save_model(self, model_type: str = 'pytorch', output_file: str = 'default'): + import os + from PIL import Image + from torch.onnx import export as onnx_export + + if output_file == 'default': + output_file = str(Path(__file__).parent) + output_file = os.path.join(output_file, 'saved', model_type) + os.makedirs(output_file, exist_ok=True) + name = self.model_name.replace('/', '-') + output_file = os.path.join(output_file, name) + if model_type in ['pytorch', 'torchscript']: + output_file = output_file + '.pt' + elif model_type == 'onnx': + output_file = output_file + '.onnx' + else: + raise AttributeError('Unsupported model_type.') + if self.modality == 'image': + sz = self.processor.feature_extractor.crop_size + dummy_input = Image.new('RGB', (sz, sz), color = 'red') + inputs = self.processor(images=dummy_input, return_tensors='pt') # a dictionary + elif self.modality == 'text': + dummy_input = 'dummy' + inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt') # a dictionary + else: + raise ValueError("modality[{}] not implemented.".format(self.modality)) + + if model_type == 'pytorch': + torch.save(self._model, output_file) + elif model_type == 'torchscript': + inputs = list(inputs.values()) + try: + try: + jit_model = torch.jit.script(self._model) + except Exception: + jit_model = torch.jit.trace(self._model, inputs, strict=False) + torch.jit.save(jit_model, output_file) + except Exception as e: + log.error(f'Fail to save as torchscript: {e}.') + raise RuntimeError(f'Fail to save as torchscript: {e}.') + elif model_type == 'onnx': + if self.modality == 'image': + input_names= ['pixel_values'] + output_names=['image_embeds'] + dynamic_axes={'pixel_values': {0: 'batch'}, 'image_embeds': {0: 'batch'}} + elif self.modality == 'text': + input_names= ['input_ids', 'attention_mask'] + output_names=['text_embeds'] + dynamic_axes={'input_ids': {0: 'batch', 1: 'sequence'}, 'attention_mask': {0: 'batch', 1: 'sequence'}, 'text_embeds': {0: 'batch'}} + else: + raise ValueError("modality[{}] not implemented.".format(self.modality)) + + onnx_export(self.model, + (dict(inputs),), + f=Path(output_file), + input_names= input_names, + output_names=output_names, + dynamic_axes=dynamic_axes, + do_constant_folding=True, + opset_version=14, + ) + else: + pass + raise NotImplementedError + diff --git a/requirements.txt b/requirements.txt index 7a0476f..91f8f77 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ torchvision torch towhee -towhee.models \ No newline at end of file +towhee.models +transformers