From debc55fc4d4a188b1c4a884199c76cb7e82223f6 Mon Sep 17 00:00:00 2001 From: mkozlowskiAzimuthe Date: Tue, 14 Mar 2023 15:48:49 +0100 Subject: [PATCH] stream trainning --- config-train.yaml | 6 +- train_stream.py | 41 ++++++++---- utils/donut_dataset.py | 4 +- utils/donut_dataset_stream.py | 113 ++++++++++++++------------------- utils/donut_model_pl_stream.py | 98 ++++++++++++++++++++++++++++ 5 files changed, 180 insertions(+), 82 deletions(-) create mode 100644 utils/donut_model_pl_stream.py diff --git a/config-train.yaml b/config-train.yaml index 67f01c4..d3857e3 100644 --- a/config-train.yaml +++ b/config-train.yaml @@ -1,14 +1,14 @@ -dataset_path: "Zombely/wikisource-small" +dataset_path: "Zombely/wikisource-yellow" pretrained_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2" start_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2" output_model_path: "Zombely/pl-donut" wandb_test_name: "wikisource-small" checkpoint_path: "./checkpoint" max_length: 768 -image_size: [2560, 1920] +image_size: [1280, 960] train_config: max_epochs: 1 - val_check_interval: 0.3 + val_check_interval: 1.0 check_val_every_n_epoch: 1 gradient_clip_val: 1.0 num_training_samples_per_epoch: 800 diff --git a/train_stream.py b/train_stream.py index 5aed52a..c53100e 100644 --- a/train_stream.py +++ b/train_stream.py @@ -9,12 +9,12 @@ from huggingface_hub import login import argparse from sconf import Config from utils.checkpoint import CustomCheckpointIO -from utils.donut_dataset_stream import DonutDataset -from utils.donut_model_pl import DonutModelPLModule +from utils.donut_dataset_stream import DonutDatasetStream +from utils.donut_model_pl_stream import DonutModelPLModuleStream from utils.callbacks import PushToHubCallback import warnings -from datasets import load_dataset - +from datasets import load_dataset, interleave_datasets +from torchdata.datapipes.iter import IterableWrapper @@ -34,8 +34,7 @@ def main(config, hug_token): added_tokens = [] - train_dataset = DonutDataset( - config.dataset_path, + train_dataset_process = DonutDatasetStream( processor=processor, model=model, max_length=config.max_length, @@ -46,8 +45,7 @@ def main(config, hug_token): sort_json_key=False, # cord dataset is preprocessed, so no need for this ) - val_dataset = DonutDataset( - config.dataset_path, + val_dataset_process = DonutDatasetStream( processor=processor, model=model, max_length=config.max_length, @@ -57,19 +55,38 @@ def main(config, hug_token): added_tokens=added_tokens, sort_json_key=False, # cord dataset is preprocessed, so no need for this ) + + dataset = load_dataset(config.dataset_path, streaming=True) + val_dataset = dataset.pop('validation') + train_dataset = interleave_datasets(list(dataset.values())) + # train_length = sum(split.num_examples for split in dataset[list(dataset.keys())[0]].info.splits.values() if split.name != 'validation') + # val_length = list(val_dataset.info.splits.values())[-1].num_examples + + + train_dataset = train_dataset.map(lambda x: train_dataset_process.process(x), remove_columns = ['image', 'ground_truth']) + val_dataset = val_dataset.map(lambda x: val_dataset_process.process(x), remove_columns = ['image', 'ground_truth']) + + # train_dataset = train_dataset.with_format('torch') + # val_dataset = val_dataset.with_format('torch') + + train_dataset = IterableWrapper(train_dataset) + val_dataset = IterableWrapper(val_dataset) model.config.pad_token_id = processor.tokenizer.pad_token_id model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([''])[0] - train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=1) - val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1) + train_dataloader = DataLoader(train_dataset, batch_size=1, num_workers=0) + val_dataloader = DataLoader(val_dataset, batch_size=1, num_workers=0) login(hug_token, True) - model_module = DonutModelPLModule(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader) + model_module = DonutModelPLModuleStream(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader) wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name) + if not os.path.exists(config.checkpoint_path): + os.mkdir(config.checkpoint_path) + checkpoint_callback = ModelCheckpoint( monitor="val_metric", dirpath=config.checkpoint_path, @@ -105,7 +122,7 @@ if __name__ == "__main__": config = Config(args.config) config.argv_update(left_argv) - hug_token = os.environ.get("HUG_TOKEN", None) + hug_token = os.environ.get("HUG_TOKEN", "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") if not torch.cuda.is_available(): warnings.warn("You don't have cuda available, training might be taking long time or impossible") diff --git a/utils/donut_dataset.py b/utils/donut_dataset.py index 200070c..dc40e6c 100644 --- a/utils/donut_dataset.py +++ b/utils/donut_dataset.py @@ -116,8 +116,8 @@ class DonutDataset(Dataset): self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer)) self.added_tokens.extend(list_of_tokens) - # def __len__(self) -> int: - # return self.dataset_length + def __len__(self) -> int: + return self.dataset_length def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ diff --git a/utils/donut_dataset_stream.py b/utils/donut_dataset_stream.py index e10a0fa..3757a4d 100644 --- a/utils/donut_dataset_stream.py +++ b/utils/donut_dataset_stream.py @@ -7,7 +7,7 @@ import torch from transformers import DonutProcessor, VisionEncoderDecoderModel -class DonutDataset(Dataset): +class DonutDatasetStream: """ DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets) Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt), @@ -24,12 +24,11 @@ class DonutDataset(Dataset): def __init__( self, - dataset_name_or_path: str, - max_length: int, processor: DonutProcessor, model: VisionEncoderDecoderModel, - split: str = "train", + max_length: int, ignore_id: int = -100, + split: str = 'train', task_start_token: str = "", prompt_end_token: str = None, sort_json_key: bool = True, @@ -37,8 +36,8 @@ class DonutDataset(Dataset): ): super().__init__() - self.max_length = max_length self.split = split + self.max_length = max_length self.processor = processor self.model = model self.ignore_id = ignore_id @@ -47,35 +46,56 @@ class DonutDataset(Dataset): self.sort_json_key = sort_json_key self.added_tokens = added_tokens - self.dataset = load_dataset(dataset_name_or_path, split=self.split, streaming=True).with_format("torch") - print(self.dataset) - self.dataset_length = len(self.dataset) + def process(self, row): - self.gt_token_sequences = [] - for sample in self.dataset: - ground_truth = json.loads(sample["ground_truth"]) - if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa - assert isinstance(ground_truth["gt_parses"], list) - gt_jsons = ground_truth["gt_parses"] - else: - assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict) - gt_jsons = [ground_truth["gt_parse"]] - self.gt_token_sequences.append( - [ - self.json2token( - gt_json, - update_special_tokens_for_json_key=self.split == "train", - sort_json_key=self.sort_json_key, - ) - + self.processor.tokenizer.eos_token - for gt_json in gt_jsons # load json from list of json - ] - ) + ground_truth = json.loads(row["ground_truth"]) + if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa + assert isinstance(ground_truth["gt_parses"], list) + gt_jsons = ground_truth["gt_parses"] + else: + assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict) + gt_jsons = [ground_truth["gt_parse"]] + + self.gt_token_sequences = ( + [ + self.json2token( + gt_json, + update_special_tokens_for_json_key=self.split == "train", + sort_json_key=self.sort_json_key, + ) + + self.processor.tokenizer.eos_token + for gt_json in gt_jsons # load json from list of json + ] + ) self.add_tokens([self.task_start_token, self.prompt_end_token]) self.prompt_end_token_id = self.processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token) + + # change if not 3 channels + if row['image'].mode != "RGB": + row['image'] = row['image'].convert("RGB") + + # inputs + pixel_values = self.processor(row["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values + pixel_values = pixel_values.squeeze() + + # targets + target_sequence = self.gt_token_sequences # can be more than one, e.g., DocVQA Task 1 + input_ids = self.processor.tokenizer( + target_sequence, + add_special_tokens=False, + max_length=self.max_length, + padding="max_length", + truncation=True, + return_tensors="pt", + )["input_ids"].squeeze(0) + + labels = input_ids.clone() + labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token + return {"pixel_values": pixel_values, "labels": labels, 'target_sequence': target_sequence } + def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True): """ Convert an ordered JSON object into a token sequence @@ -117,40 +137,3 @@ class DonutDataset(Dataset): self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer)) self.added_tokens.extend(list_of_tokens) - # def __len__(self) -> int: - # return self.dataset_length - - def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Load image from image_path of given dataset_path and convert into input_tensor and labels - Convert gt data into input_ids (tokenized string) - Returns: - input_tensor : preprocessed image - input_ids : tokenized gt_data - labels : masked labels (model doesn't need to predict prompt and pad token) - """ - sample = self.dataset[idx] - - # change if not 3 channels - if sample['image'].mode != "RGB": - sample['image'] = sample['image'].convert("RGB") - - # inputs - pixel_values = self.processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values - pixel_values = pixel_values.squeeze() - - # targets - target_sequence = random.choice(self.gt_token_sequences[idx]) # can be more than one, e.g., DocVQA Task 1 - input_ids = self.processor.tokenizer( - target_sequence, - add_special_tokens=False, - max_length=self.max_length, - padding="max_length", - truncation=True, - return_tensors="pt", - )["input_ids"].squeeze(0) - - labels = input_ids.clone() - labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token - # labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id # model doesn't need to predict prompt (for VQA) - return pixel_values, labels, target_sequence diff --git a/utils/donut_model_pl_stream.py b/utils/donut_model_pl_stream.py new file mode 100644 index 0000000..79fd91d --- /dev/null +++ b/utils/donut_model_pl_stream.py @@ -0,0 +1,98 @@ +import torch +import pytorch_lightning as pl +from nltk import edit_distance +import re +import numpy as np + + +class DonutModelPLModuleStream(pl.LightningModule): + def __init__(self, config, processor, model, max_length, train_dataloader, val_dataloader): + super().__init__() + self.config = config + self.processor = processor + self.model = model + self.max_length = max_length + self._train_dataloader = train_dataloader + self._val_dataloader = val_dataloader + + def training_step(self, batch, batch_idx): + # pixel_values, labels, _ = batch + pixel_values = batch['pixel_values'] + labels = batch['labels'] + outputs = self.model(pixel_values, labels=labels) + loss = outputs.loss + self.log_dict({"train_loss": loss}, sync_dist=True) + return loss + + def validation_step(self, batch, batch_idx, dataset_idx=0): + # pixel_values, labels, answers = batch + + pixel_values = batch['pixel_values'] + labels = batch['labels'] + answers = batch['target_sequence'] + batch_size = pixel_values.shape[0] + # we feed the prompt to the model + decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device) + + outputs = self.model.generate(pixel_values, + decoder_input_ids=decoder_input_ids, + max_length=self.max_length, + early_stopping=True, + pad_token_id=self.processor.tokenizer.pad_token_id, + eos_token_id=self.processor.tokenizer.eos_token_id, + use_cache=True, + num_beams=1, + bad_words_ids=[[self.processor.tokenizer.unk_token_id]], + return_dict_in_generate=True,) + + predictions = [] + for seq in self.processor.tokenizer.batch_decode(outputs.sequences): + seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "") + seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token + predictions.append(seq) + + scores = list() + for pred, answer in zip(predictions, answers): + pred = re.sub(r"(?:(?<=>) | (?=", "", answer, count=1) + answer = answer.replace(self.processor.tokenizer.eos_token, "") + scores.append(edit_distance(pred, answer) / max(len(pred), len(answer))) + + if self.config.get("verbose", False) and len(scores) == 1: + print(f"Prediction: {pred}") + print(f" Answer: {answer}") + print(f" Normed ED: {scores[0]}") + + return scores + + def validation_epoch_end(self, validation_step_outputs): + # I set this to 1 manually + # (previously set to len(self.config.dataset_name_or_paths)) + num_of_loaders = 1 + if num_of_loaders == 1: + validation_step_outputs = [validation_step_outputs] + assert len(validation_step_outputs) == num_of_loaders + cnt = [0] * num_of_loaders + total_metric = [0] * num_of_loaders + val_metric = [0] * num_of_loaders + for i, results in enumerate(validation_step_outputs): + for scores in results: + cnt[i] += len(scores) + total_metric[i] += np.sum(scores) + val_metric[i] = total_metric[i] / cnt[i] + val_metric_name = f"val_metric_{i}th_dataset" + self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True) + self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True) + + def configure_optimizers(self): + # TODO add scheduler + optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr")) + + return optimizer + + def train_dataloader(self): + return self._train_dataloader + + def val_dataloader(self): + return self._val_dataloader \ No newline at end of file