From 67f4c12f2de09d0593bfd5e1802007bc59a6b8e0 Mon Sep 17 00:00:00 2001 From: jinlingxu06 Date: Fri, 24 Feb 2023 18:31:46 +0800 Subject: [PATCH] add training for blip. Signed-off-by: jinlingxu06 --- blip.py | 107 ++++++++- train_blip_with_hf_trainer.py | 412 ++++++++++++++++++++++++++++++++++ train_clip_with_hf_trainer.py | 412 ++++++++++++++++++++++++++++++++++ 3 files changed, 930 insertions(+), 1 deletion(-) create mode 100644 train_blip_with_hf_trainer.py create mode 100644 train_clip_with_hf_trainer.py diff --git a/blip.py b/blip.py index cc88119..fb503ce 100644 --- a/blip.py +++ b/blip.py @@ -15,6 +15,7 @@ import sys import os from pathlib import Path +from typing import Optional, Union, Tuple import torch import logging @@ -23,6 +24,8 @@ from torch import nn from transformers import AutoProcessor, BlipForImageTextRetrieval from transformers import logging as t_logging +from transformers.models.blip.modeling_blip import BlipOutput, blip_loss + from towhee import register from towhee.operator.base import NNOperator, OperatorFlag from towhee.types.arg import arg, to_image_color @@ -33,6 +36,99 @@ warnings.filterwarnings('ignore') os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' t_logging.set_verbosity_error() +class BlipForImageTextRetrievalFinetune(): + def __init__(self, blip_itm_model): + self.proxied = blip_itm_model + self.proxied.logit_scale = torch.nn.Parameter(torch.ones([]) * blip_itm_model.config.logit_scale_init_value) + + def __getattr__(self, attr): + if attr == 'proxied': + return self.proxied + elif attr != 'forward': + if getattr(self.proxied, attr) is True: + def wrapped_method(*args, **kwargs): + result = getattr(self.proxied, attr)(*args, **kwargs) + return result + return wrapped_method + else: + return getattr(self.proxied, attr) + else: + return self._forward + + def forward(self, *args, **kwargs): + raise NotImplemented + + def __call__(self, *args, **kwargs): + return self._forward(*args, **kwargs) + + def _forward( + self, + input_ids: Optional[torch.LongTensor] = None, + pixel_values: Optional[torch.FloatTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + return_loss: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BlipOutput]: + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + vision_outputs = self.vision_model( + pixel_values=pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + text_outputs = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + image_embeds = vision_outputs[1] + image_embeds = self.vision_proj(image_embeds) + + text_embeds = text_outputs[0] + text_embeds = self.text_proj(text_embeds[:,0,:]) + + # normalized features + image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True) + text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True) + + # cosine similarity as logits + logit_scale = self.logit_scale.exp() + logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale + logits_per_image = logits_per_text.t() + + loss = None + if return_loss: + loss = blip_loss(logits_per_text) + + if not return_dict: + output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs) + return ((loss,) + output) if loss is not None else output + + return BlipOutput( + loss=loss, + logits_per_image=logits_per_image, + logits_per_text=logits_per_text, + text_embeds=text_embeds, + image_embeds=image_embeds, + text_model_output=text_outputs, + vision_model_output=vision_outputs, + ) + + def create_model(cfg, modality, checkpoint_path, device): hf_blip_model = BlipForImageTextRetrieval.from_pretrained(cfg) if checkpoint_path: @@ -158,7 +254,16 @@ class Blip(NNOperator): return self.model.model def train(self, **kwargs): - raise NotImplementedError + import sys + import pathlib + path = str(pathlib.Path(__file__).parent) + sys.path.append(path) + from train_blip_with_hf_trainer import train_with_hf_trainer + data_args = kwargs.pop('data_args', None) + training_args = kwargs.pop('training_args', None) + model_finetune = BlipForImageTextRetrievalFinetune(self._model.backbone) + train_with_hf_trainer(model_finetune, self.processor.tokenizer, data_args, training_args) + @property def supported_formats(self): diff --git a/train_blip_with_hf_trainer.py b/train_blip_with_hf_trainer.py new file mode 100644 index 0000000..f1c26d7 --- /dev/null +++ b/train_blip_with_hf_trainer.py @@ -0,0 +1,412 @@ +import logging +import os +import sys +import transformers +import dataclasses +from dataclasses import dataclass, field +from typing import Optional, List + +import torch +from datasets import load_dataset +from PIL import Image +from torchvision.io import ImageReadMode, read_image +from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize +from torchvision.transforms.functional import InterpolationMode + +from transformers import ( + MODEL_FOR_CAUSAL_LM_MAPPING, + TrainingArguments, + default_data_collator, + is_torch_tpu_available, + set_seed, +) +from transformers.trainer_utils import get_last_checkpoint + +# We use torchvision for faster image pre-processing. The transforms are implemented as nn.Module, +# so we jit it to be faster. + + +logger = logging.getLogger(__name__) + +dataset_name_mapping = { + "image_caption_dataset.py": ("image_path", "caption"), +} + + +def dataclass_from_dict(klass, d): + try: + fieldtypes = {f.name: f.type for f in dataclasses.fields(klass)} + return klass(**{f: dataclass_from_dict(fieldtypes[f], d[f]) for f in d}) + except: + return d # Not a dataclass field + + +@dataclass +class DataTrainingArguments: + """ + Arguments pertaining to what data we are going to input our model for training and eval. + """ + + dataset_name: Optional[str] = field( + default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} + ) + dataset_config_name: Optional[str] = field( + default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} + ) + data_dir: Optional[str] = field(default=None, metadata={"help": "The data directory containing input files."}) + image_column: Optional[str] = field( + default="image_path", + metadata={"help": "The name of the column in the datasets containing the full image file paths."}, + ) + caption_column: Optional[str] = field( + default="caption", + metadata={"help": "The name of the column in the datasets containing the image captions."}, + ) + train_file: Optional[str] = field( + default=None, metadata={"help": "The input training data file (a jsonlines file)."} + ) + validation_file: Optional[str] = field( + default=None, + metadata={"help": "An optional input evaluation data file (a jsonlines file)."}, + ) + max_seq_length: Optional[int] = field( + default=77, + metadata={ + "help": ( + "The maximum total input sequence length after tokenization. Sequences longer " + "than this will be truncated, sequences shorter will be padded." + ) + }, + ) + max_train_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ) + }, + ) + max_eval_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of evaluation examples to this " + "value if set." + ) + }, + ) + overwrite_cache: bool = field( + default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} + ) + preprocessing_num_workers: Optional[int] = field( + default=None, + metadata={"help": "The number of processes to use for the preprocessing."}, + ) + cache_dir: Optional[str] = field( + default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"} + ) + image_mean: Optional[str] = field( + default=None, metadata={"help": "image preprocessing mean"} + ) + image_std: Optional[str] = field( + default=None, metadata={"help": "image preprocessing std"} + ) + freeze_vision_model: bool = field( + default=False, metadata={"help":"Whether to freeze the vision model parameters or not."} + ) + freeze_text_model: bool = field( + default=False, metadata={"help": "Whether to freeze the text model parameters or not."} + ) + + + def __post_init__(self): + if self.dataset_name is None and self.train_file is None and self.validation_file is None: + raise ValueError("Need either a dataset name or a training/validation file.") + else: + if self.train_file is not None: + extension = self.train_file.split(".")[-1] + assert extension in ["csv", "json"], "`train_file` should be a csv or a json file." + if self.validation_file is not None: + extension = self.validation_file.split(".")[-1] + assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." + if self.validation_file is not None: + extension = self.validation_file.split(".")[-1] + assert extension == "json", "`validation_file` should be a json file." + + +class Transform(torch.nn.Module): + def __init__(self, image_size, mean, std): + super().__init__() + self.transforms = torch.nn.Sequential( + Resize([image_size], interpolation=InterpolationMode.BICUBIC), + CenterCrop(image_size), + ConvertImageDtype(torch.float), + Normalize(mean, std), + ) + + def forward(self, x) -> torch.Tensor: + """`x` should be an instance of `PIL.Image.Image`""" + with torch.no_grad(): + x = self.transforms(x) + return x + +def collate_fn(examples): + pixel_values = torch.stack([example["pixel_values"] for example in examples]) + input_ids = torch.tensor([example["input_ids"] for example in examples], dtype=torch.long) + attention_mask = torch.tensor([example["attention_mask"] for example in examples], dtype=torch.long) + return { + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + "return_loss": True, + } + + +def train_with_hf_trainer(model, tokenizer, data_args, training_args, **kwargs): + + import evaluate + import datasets + + from transformers import Trainer + from datasets import load_dataset + from towhee.trainer.training_config import get_dataclasses_help + + print('**** DataTrainingArguments ****') + get_dataclasses_help(DataTrainingArguments) + data_args = dataclass_from_dict(DataTrainingArguments, data_args) + + print('**** TrainingArguments ****') + get_dataclasses_help(TrainingArguments) + training_args = dataclass_from_dict(TrainingArguments, training_args) + + # Setup logging + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], + ) + + # Setup logging + #+ training_args + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + temp_cache_dir = data_args.cache_dir + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" + ) + logger.info(f"Training/evaluation parameters {training_args}") + + # Detecting last checkpoint and eventualy continue from last checkpoint + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: + raise ValueError( + f"Output directory ({training_args.output_dir}) already exists and is not empty. " + "Use --overwrite_output_dir to overcome." + ) + elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: + logger.info( + f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + + # Load dataset + # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below) + # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/ + # (the dataset will be downloaded automatically from the datasets Hub). + # + # For CSV/JSON files this script will use the first column for the full image path and the second column for the + # captions (unless you specify column names for this with the `image_column` and `caption_column` arguments). + # + + if data_args.dataset_name is not None: + # Downloading and loading a dataset from the hub. + dataset = load_dataset( + data_args.dataset_name, + data_args.dataset_config_name, + cache_dir=temp_cache_dir, + keep_in_memory=False, + data_dir=data_args.data_dir, + # use_auth_token=True if model_args.use_auth_token else None, + ) + else: + data_files = {} + if data_args.train_file is not None: + data_files["train"] = data_args.train_file + extension = data_args.train_file.split(".")[-1] + if data_args.validation_file is not None: + data_files["validation"] = data_args.validation_file + extension = data_args.validation_file.split(".")[-1] + dataset = load_dataset( + extension, + data_files=data_files, + cache_dir=temp_cache_dir, + # use_auth_token=True if model_args.use_auth_token else None, + ) + + config = model.config + + freeze_vision_model = data_args.freeze_vision_model + freeze_text_model = data_args.freeze_text_model + + def _freeze_params(module): + for param in module.parameters(): + param.requires_grad = False + + if model_args.freeze_vision_model: + _freeze_params(model.vision_model) + + if model_args.freeze_text_model: + _freeze_params(model.text_model) + + if freeze_vision_model is True: + _freeze_params(model.vision_model) + + if freeze_text_model is True: + _freeze_params(model.text_model) + + set_seed(training_args.seed) + + if training_args.do_train: + column_names = dataset["train"].column_names + elif training_args.do_eval: + column_names = dataset["validation"].column_names + else: + logger.info("There is nothing to do. Please pass `do_train`, `do_eval`.") + return + + dataset_columns = dataset_name_mapping.get(data_args.dataset_name, None) + if data_args.image_column is None: + image_column = dataset_columns[0] if dataset_columns is not None else column_names[0] + else: + image_column = data_args.image_column + if image_column not in column_names: + raise ValueError( + f"--image_column' value '{data_args.image_column}' needs to be one of: {', '.join(column_names)}" + ) + if data_args.caption_column is None: + caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1] + else: + caption_column = data_args.caption_column + if caption_column not in column_names: + raise ValueError( + f"--caption_column' value '{data_args.caption_column}' needs to be one of: {', '.join(column_names)}" + ) + + + image_mean, image_std = data_args.image_mean, data_args.image_std + # Preprocessing the datasets. + # Initialize torchvision transforms and jit it for faster processing. + image_transformations = Transform( + config.vision_config.image_size, image_mean, image_std + ) + image_transformations = torch.jit.script(image_transformations) + + # Preprocessing the datasets. + # We need to tokenize input captions and transform the images. + #data_args + + def tokenize_captions(examples): + captions = [caption for caption in examples[caption_column]] + text_inputs = tokenizer(captions, max_length=data_args.max_seq_length, padding="max_length", truncation=True) + examples["input_ids"] = text_inputs.input_ids + examples["attention_mask"] = text_inputs.attention_mask + return examples + + def transform_images(examples): + images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[image_column]] + examples["pixel_values"] = [image_transformations(image) for image in images] + return examples + + def filter_corrupt_images(examples): + """remove problematic images""" + valid_images = [] + for image_file in examples[image_column]: + try: + Image.open(image_file) + valid_images.append(True) + except Exception: + valid_images.append(False) + return valid_images + + if training_args.do_train: + if "train" not in dataset: + raise ValueError("--do_train requires a train dataset") + train_dataset = dataset["train"] + if data_args.max_train_samples is not None: + max_train_samples = min(len(train_dataset), data_args.max_train_samples) + train_dataset = train_dataset.select(range(max_train_samples)) + train_dataset = train_dataset.filter( + filter_corrupt_images, batched=True, num_proc=data_args.preprocessing_num_workers + ) + train_dataset = train_dataset.map( + function=tokenize_captions, + batched=True, + remove_columns=[col for col in column_names if col != image_column], + num_proc=data_args.preprocessing_num_workers, + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on train dataset", + ) + + # Transform images on the fly as doing it on the whole dataset takes too much time. + train_dataset.set_transform(transform_images) + + if training_args.do_eval: + if "validation" not in dataset: + raise ValueError("--do_eval requires a train validation") + eval_dataset = dataset["validation"] + if data_args.max_eval_samples is not None: + max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) + eval_dataset = eval_dataset.select(range(max_eval_samples)) + + eval_dataset = eval_dataset.filter( + filter_corrupt_images, batched=True, num_proc=data_args.preprocessing_num_workers + ) + eval_dataset = eval_dataset.map( + function=tokenize_captions, + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=[col for col in column_names if col != image_column], + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on validation dataset", + ) + + # Transform images on the fly as doing it on the whole dataset takes too much time. + eval_dataset.set_transform(transform_images) + + # Initalize our trainer + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset if training_args.do_train else None, + eval_dataset=eval_dataset if training_args.do_eval else None, + data_collator=collate_fn, + ) + + # Training + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + train_result = trainer.train(resume_from_checkpoint=checkpoint) + trainer.save_model() + trainer.log_metrics("train", train_result.metrics) + trainer.save_metrics("train", train_result.metrics) + trainer.save_state() + # Evaluation + if training_args.do_eval: + metrics = trainer.evaluate() + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + + diff --git a/train_clip_with_hf_trainer.py b/train_clip_with_hf_trainer.py new file mode 100644 index 0000000..f1c26d7 --- /dev/null +++ b/train_clip_with_hf_trainer.py @@ -0,0 +1,412 @@ +import logging +import os +import sys +import transformers +import dataclasses +from dataclasses import dataclass, field +from typing import Optional, List + +import torch +from datasets import load_dataset +from PIL import Image +from torchvision.io import ImageReadMode, read_image +from torchvision.transforms import CenterCrop, ConvertImageDtype, Normalize, Resize +from torchvision.transforms.functional import InterpolationMode + +from transformers import ( + MODEL_FOR_CAUSAL_LM_MAPPING, + TrainingArguments, + default_data_collator, + is_torch_tpu_available, + set_seed, +) +from transformers.trainer_utils import get_last_checkpoint + +# We use torchvision for faster image pre-processing. The transforms are implemented as nn.Module, +# so we jit it to be faster. + + +logger = logging.getLogger(__name__) + +dataset_name_mapping = { + "image_caption_dataset.py": ("image_path", "caption"), +} + + +def dataclass_from_dict(klass, d): + try: + fieldtypes = {f.name: f.type for f in dataclasses.fields(klass)} + return klass(**{f: dataclass_from_dict(fieldtypes[f], d[f]) for f in d}) + except: + return d # Not a dataclass field + + +@dataclass +class DataTrainingArguments: + """ + Arguments pertaining to what data we are going to input our model for training and eval. + """ + + dataset_name: Optional[str] = field( + default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."} + ) + dataset_config_name: Optional[str] = field( + default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} + ) + data_dir: Optional[str] = field(default=None, metadata={"help": "The data directory containing input files."}) + image_column: Optional[str] = field( + default="image_path", + metadata={"help": "The name of the column in the datasets containing the full image file paths."}, + ) + caption_column: Optional[str] = field( + default="caption", + metadata={"help": "The name of the column in the datasets containing the image captions."}, + ) + train_file: Optional[str] = field( + default=None, metadata={"help": "The input training data file (a jsonlines file)."} + ) + validation_file: Optional[str] = field( + default=None, + metadata={"help": "An optional input evaluation data file (a jsonlines file)."}, + ) + max_seq_length: Optional[int] = field( + default=77, + metadata={ + "help": ( + "The maximum total input sequence length after tokenization. Sequences longer " + "than this will be truncated, sequences shorter will be padded." + ) + }, + ) + max_train_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of training examples to this " + "value if set." + ) + }, + ) + max_eval_samples: Optional[int] = field( + default=None, + metadata={ + "help": ( + "For debugging purposes or quicker training, truncate the number of evaluation examples to this " + "value if set." + ) + }, + ) + overwrite_cache: bool = field( + default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} + ) + preprocessing_num_workers: Optional[int] = field( + default=None, + metadata={"help": "The number of processes to use for the preprocessing."}, + ) + cache_dir: Optional[str] = field( + default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"} + ) + image_mean: Optional[str] = field( + default=None, metadata={"help": "image preprocessing mean"} + ) + image_std: Optional[str] = field( + default=None, metadata={"help": "image preprocessing std"} + ) + freeze_vision_model: bool = field( + default=False, metadata={"help":"Whether to freeze the vision model parameters or not."} + ) + freeze_text_model: bool = field( + default=False, metadata={"help": "Whether to freeze the text model parameters or not."} + ) + + + def __post_init__(self): + if self.dataset_name is None and self.train_file is None and self.validation_file is None: + raise ValueError("Need either a dataset name or a training/validation file.") + else: + if self.train_file is not None: + extension = self.train_file.split(".")[-1] + assert extension in ["csv", "json"], "`train_file` should be a csv or a json file." + if self.validation_file is not None: + extension = self.validation_file.split(".")[-1] + assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." + if self.validation_file is not None: + extension = self.validation_file.split(".")[-1] + assert extension == "json", "`validation_file` should be a json file." + + +class Transform(torch.nn.Module): + def __init__(self, image_size, mean, std): + super().__init__() + self.transforms = torch.nn.Sequential( + Resize([image_size], interpolation=InterpolationMode.BICUBIC), + CenterCrop(image_size), + ConvertImageDtype(torch.float), + Normalize(mean, std), + ) + + def forward(self, x) -> torch.Tensor: + """`x` should be an instance of `PIL.Image.Image`""" + with torch.no_grad(): + x = self.transforms(x) + return x + +def collate_fn(examples): + pixel_values = torch.stack([example["pixel_values"] for example in examples]) + input_ids = torch.tensor([example["input_ids"] for example in examples], dtype=torch.long) + attention_mask = torch.tensor([example["attention_mask"] for example in examples], dtype=torch.long) + return { + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + "return_loss": True, + } + + +def train_with_hf_trainer(model, tokenizer, data_args, training_args, **kwargs): + + import evaluate + import datasets + + from transformers import Trainer + from datasets import load_dataset + from towhee.trainer.training_config import get_dataclasses_help + + print('**** DataTrainingArguments ****') + get_dataclasses_help(DataTrainingArguments) + data_args = dataclass_from_dict(DataTrainingArguments, data_args) + + print('**** TrainingArguments ****') + get_dataclasses_help(TrainingArguments) + training_args = dataclass_from_dict(TrainingArguments, training_args) + + # Setup logging + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + handlers=[logging.StreamHandler(sys.stdout)], + ) + + # Setup logging + #+ training_args + log_level = training_args.get_process_log_level() + logger.setLevel(log_level) + transformers.utils.logging.set_verbosity(log_level) + transformers.utils.logging.enable_default_handler() + transformers.utils.logging.enable_explicit_format() + + temp_cache_dir = data_args.cache_dir + logger.warning( + f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" + + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" + ) + logger.info(f"Training/evaluation parameters {training_args}") + + # Detecting last checkpoint and eventualy continue from last checkpoint + last_checkpoint = None + if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: + last_checkpoint = get_last_checkpoint(training_args.output_dir) + if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: + raise ValueError( + f"Output directory ({training_args.output_dir}) already exists and is not empty. " + "Use --overwrite_output_dir to overcome." + ) + elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: + logger.info( + f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " + "the `--output_dir` or add `--overwrite_output_dir` to train from scratch." + ) + + + # Load dataset + # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below) + # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/ + # (the dataset will be downloaded automatically from the datasets Hub). + # + # For CSV/JSON files this script will use the first column for the full image path and the second column for the + # captions (unless you specify column names for this with the `image_column` and `caption_column` arguments). + # + + if data_args.dataset_name is not None: + # Downloading and loading a dataset from the hub. + dataset = load_dataset( + data_args.dataset_name, + data_args.dataset_config_name, + cache_dir=temp_cache_dir, + keep_in_memory=False, + data_dir=data_args.data_dir, + # use_auth_token=True if model_args.use_auth_token else None, + ) + else: + data_files = {} + if data_args.train_file is not None: + data_files["train"] = data_args.train_file + extension = data_args.train_file.split(".")[-1] + if data_args.validation_file is not None: + data_files["validation"] = data_args.validation_file + extension = data_args.validation_file.split(".")[-1] + dataset = load_dataset( + extension, + data_files=data_files, + cache_dir=temp_cache_dir, + # use_auth_token=True if model_args.use_auth_token else None, + ) + + config = model.config + + freeze_vision_model = data_args.freeze_vision_model + freeze_text_model = data_args.freeze_text_model + + def _freeze_params(module): + for param in module.parameters(): + param.requires_grad = False + + if model_args.freeze_vision_model: + _freeze_params(model.vision_model) + + if model_args.freeze_text_model: + _freeze_params(model.text_model) + + if freeze_vision_model is True: + _freeze_params(model.vision_model) + + if freeze_text_model is True: + _freeze_params(model.text_model) + + set_seed(training_args.seed) + + if training_args.do_train: + column_names = dataset["train"].column_names + elif training_args.do_eval: + column_names = dataset["validation"].column_names + else: + logger.info("There is nothing to do. Please pass `do_train`, `do_eval`.") + return + + dataset_columns = dataset_name_mapping.get(data_args.dataset_name, None) + if data_args.image_column is None: + image_column = dataset_columns[0] if dataset_columns is not None else column_names[0] + else: + image_column = data_args.image_column + if image_column not in column_names: + raise ValueError( + f"--image_column' value '{data_args.image_column}' needs to be one of: {', '.join(column_names)}" + ) + if data_args.caption_column is None: + caption_column = dataset_columns[1] if dataset_columns is not None else column_names[1] + else: + caption_column = data_args.caption_column + if caption_column not in column_names: + raise ValueError( + f"--caption_column' value '{data_args.caption_column}' needs to be one of: {', '.join(column_names)}" + ) + + + image_mean, image_std = data_args.image_mean, data_args.image_std + # Preprocessing the datasets. + # Initialize torchvision transforms and jit it for faster processing. + image_transformations = Transform( + config.vision_config.image_size, image_mean, image_std + ) + image_transformations = torch.jit.script(image_transformations) + + # Preprocessing the datasets. + # We need to tokenize input captions and transform the images. + #data_args + + def tokenize_captions(examples): + captions = [caption for caption in examples[caption_column]] + text_inputs = tokenizer(captions, max_length=data_args.max_seq_length, padding="max_length", truncation=True) + examples["input_ids"] = text_inputs.input_ids + examples["attention_mask"] = text_inputs.attention_mask + return examples + + def transform_images(examples): + images = [read_image(image_file, mode=ImageReadMode.RGB) for image_file in examples[image_column]] + examples["pixel_values"] = [image_transformations(image) for image in images] + return examples + + def filter_corrupt_images(examples): + """remove problematic images""" + valid_images = [] + for image_file in examples[image_column]: + try: + Image.open(image_file) + valid_images.append(True) + except Exception: + valid_images.append(False) + return valid_images + + if training_args.do_train: + if "train" not in dataset: + raise ValueError("--do_train requires a train dataset") + train_dataset = dataset["train"] + if data_args.max_train_samples is not None: + max_train_samples = min(len(train_dataset), data_args.max_train_samples) + train_dataset = train_dataset.select(range(max_train_samples)) + train_dataset = train_dataset.filter( + filter_corrupt_images, batched=True, num_proc=data_args.preprocessing_num_workers + ) + train_dataset = train_dataset.map( + function=tokenize_captions, + batched=True, + remove_columns=[col for col in column_names if col != image_column], + num_proc=data_args.preprocessing_num_workers, + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on train dataset", + ) + + # Transform images on the fly as doing it on the whole dataset takes too much time. + train_dataset.set_transform(transform_images) + + if training_args.do_eval: + if "validation" not in dataset: + raise ValueError("--do_eval requires a train validation") + eval_dataset = dataset["validation"] + if data_args.max_eval_samples is not None: + max_eval_samples = min(len(eval_dataset), data_args.max_eval_samples) + eval_dataset = eval_dataset.select(range(max_eval_samples)) + + eval_dataset = eval_dataset.filter( + filter_corrupt_images, batched=True, num_proc=data_args.preprocessing_num_workers + ) + eval_dataset = eval_dataset.map( + function=tokenize_captions, + batched=True, + num_proc=data_args.preprocessing_num_workers, + remove_columns=[col for col in column_names if col != image_column], + load_from_cache_file=not data_args.overwrite_cache, + desc="Running tokenizer on validation dataset", + ) + + # Transform images on the fly as doing it on the whole dataset takes too much time. + eval_dataset.set_transform(transform_images) + + # Initalize our trainer + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset if training_args.do_train else None, + eval_dataset=eval_dataset if training_args.do_eval else None, + data_collator=collate_fn, + ) + + # Training + if training_args.do_train: + checkpoint = None + if training_args.resume_from_checkpoint is not None: + checkpoint = training_args.resume_from_checkpoint + elif last_checkpoint is not None: + checkpoint = last_checkpoint + train_result = trainer.train(resume_from_checkpoint=checkpoint) + trainer.save_model() + trainer.log_metrics("train", train_result.metrics) + trainer.save_metrics("train", train_result.metrics) + trainer.save_state() + # Evaluation + if training_args.do_eval: + metrics = trainer.evaluate() + trainer.log_metrics("eval", metrics) + trainer.save_metrics("eval", metrics) + +