blip
copied
jinlingxu06
2 years ago
4 changed files with 94 additions and 428 deletions
@ -1,412 +0,0 @@ |
|||
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) |
|||
|
|||
|
|||
Loading…
Reference in new issue