stream trainning

This commit is contained in:
mkozlowskiAzimuthe 2023-03-14 15:48:49 +01:00
parent 93a231a477
commit debc55fc4d
5 changed files with 180 additions and 82 deletions

View File

@ -1,14 +1,14 @@
dataset_path: "Zombely/wikisource-small" dataset_path: "Zombely/wikisource-yellow"
pretrained_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2" pretrained_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2"
start_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" output_model_path: "Zombely/pl-donut"
wandb_test_name: "wikisource-small" wandb_test_name: "wikisource-small"
checkpoint_path: "./checkpoint" checkpoint_path: "./checkpoint"
max_length: 768 max_length: 768
image_size: [2560, 1920] image_size: [1280, 960]
train_config: train_config:
max_epochs: 1 max_epochs: 1
val_check_interval: 0.3 val_check_interval: 1.0
check_val_every_n_epoch: 1 check_val_every_n_epoch: 1
gradient_clip_val: 1.0 gradient_clip_val: 1.0
num_training_samples_per_epoch: 800 num_training_samples_per_epoch: 800

View File

@ -9,12 +9,12 @@ from huggingface_hub import login
import argparse import argparse
from sconf import Config from sconf import Config
from utils.checkpoint import CustomCheckpointIO from utils.checkpoint import CustomCheckpointIO
from utils.donut_dataset_stream import DonutDataset from utils.donut_dataset_stream import DonutDatasetStream
from utils.donut_model_pl import DonutModelPLModule from utils.donut_model_pl_stream import DonutModelPLModuleStream
from utils.callbacks import PushToHubCallback from utils.callbacks import PushToHubCallback
import warnings 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 = [] added_tokens = []
train_dataset = DonutDataset( train_dataset_process = DonutDatasetStream(
config.dataset_path,
processor=processor, processor=processor,
model=model, model=model,
max_length=config.max_length, 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 sort_json_key=False, # cord dataset is preprocessed, so no need for this
) )
val_dataset = DonutDataset( val_dataset_process = DonutDatasetStream(
config.dataset_path,
processor=processor, processor=processor,
model=model, model=model,
max_length=config.max_length, max_length=config.max_length,
@ -57,19 +55,38 @@ def main(config, hug_token):
added_tokens=added_tokens, added_tokens=added_tokens,
sort_json_key=False, # cord dataset is preprocessed, so no need for this 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.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0] model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=1) train_dataloader = DataLoader(train_dataset, batch_size=1, num_workers=0)
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1) val_dataloader = DataLoader(val_dataset, batch_size=1, num_workers=0)
login(hug_token, True) 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) 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( checkpoint_callback = ModelCheckpoint(
monitor="val_metric", monitor="val_metric",
dirpath=config.checkpoint_path, dirpath=config.checkpoint_path,
@ -105,7 +122,7 @@ if __name__ == "__main__":
config = Config(args.config) config = Config(args.config)
config.argv_update(left_argv) 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(): if not torch.cuda.is_available():
warnings.warn("You don't have cuda available, training might be taking long time or impossible") warnings.warn("You don't have cuda available, training might be taking long time or impossible")

View File

@ -116,8 +116,8 @@ class DonutDataset(Dataset):
self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer)) self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer))
self.added_tokens.extend(list_of_tokens) self.added_tokens.extend(list_of_tokens)
# def __len__(self) -> int: def __len__(self) -> int:
# return self.dataset_length return self.dataset_length
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
""" """

View File

@ -7,7 +7,7 @@ import torch
from transformers import DonutProcessor, VisionEncoderDecoderModel 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) 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), Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
@ -24,12 +24,11 @@ class DonutDataset(Dataset):
def __init__( def __init__(
self, self,
dataset_name_or_path: str,
max_length: int,
processor: DonutProcessor, processor: DonutProcessor,
model: VisionEncoderDecoderModel, model: VisionEncoderDecoderModel,
split: str = "train", max_length: int,
ignore_id: int = -100, ignore_id: int = -100,
split: str = 'train',
task_start_token: str = "<s>", task_start_token: str = "<s>",
prompt_end_token: str = None, prompt_end_token: str = None,
sort_json_key: bool = True, sort_json_key: bool = True,
@ -37,8 +36,8 @@ class DonutDataset(Dataset):
): ):
super().__init__() super().__init__()
self.max_length = max_length
self.split = split self.split = split
self.max_length = max_length
self.processor = processor self.processor = processor
self.model = model self.model = model
self.ignore_id = ignore_id self.ignore_id = ignore_id
@ -47,35 +46,56 @@ class DonutDataset(Dataset):
self.sort_json_key = sort_json_key self.sort_json_key = sort_json_key
self.added_tokens = added_tokens self.added_tokens = added_tokens
self.dataset = load_dataset(dataset_name_or_path, split=self.split, streaming=True).with_format("torch") def process(self, row):
print(self.dataset)
self.dataset_length = len(self.dataset)
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( ground_truth = json.loads(row["ground_truth"])
[ if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa
self.json2token( assert isinstance(ground_truth["gt_parses"], list)
gt_json, gt_jsons = ground_truth["gt_parses"]
update_special_tokens_for_json_key=self.split == "train", else:
sort_json_key=self.sort_json_key, assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
) gt_jsons = [ground_truth["gt_parse"]]
+ self.processor.tokenizer.eos_token
for gt_json in gt_jsons # load json from list of json 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.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) 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): 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 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.model.decoder.resize_token_embeddings(len(self.processor.tokenizer))
self.added_tokens.extend(list_of_tokens) 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

View File

@ -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"(?:(?<=>) | (?=</s_))", "", pred)
# NOT NEEDED ANYMORE
# answer = 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