s444415/donut
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

teesting train2 2

Browse Source
This commit is contained in:
Michał Kozłowski 2023-01-25 21:43:13 +01:00
parent dd5febad65
commit ecce4427a5
2 changed files with 54 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

107
train_stream.py
View File

@ -35,71 +35,70 @@ def main(config, hug_token):
added_tokens = [] added_tokens = []
dataset = load_dataset(config.dataset_path) dataset = load_dataset(config.dataset_path)
dataset.train_test_split(test_size=0.1) dataset = dataset.train_test_split(test_size=0.1)
print(dataset)
# train_dataset = DonutDataset( train_dataset = DonutDataset(
# dataset, dataset,
# processor=processor, processor=processor,
# model=model, model=model,
# max_length=config.max_length, max_length=config.max_length,
# split="train", split="train",
# task_start_token="<s_cord-v2>", task_start_token="<s_cord-v2>",
# prompt_end_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
# 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
# ) )
# val_dataset = DonutDataset( val_dataset = DonutDataset(
# dataset, dataset,
# processor=processor, processor=processor,
# model=model, model=model,
# max_length=config.max_length, max_length=config.max_length,
# split="validation", split="test",
# task_start_token="<s_cord-v2>", task_start_token="<s_cord-v2>",
# prompt_end_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
# 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
# ) )
# 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, shuffle=True, num_workers=1)
# val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1) val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1)
# 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 = DonutModelPLModule(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)
# checkpoint_callback = ModelCheckpoint( checkpoint_callback = ModelCheckpoint(
# monitor="val_metric", monitor="val_metric",
# dirpath=config.checkpoint_path, dirpath=config.checkpoint_path,
# filename="artifacts", filename="artifacts",
# save_top_k=1, save_top_k=1,
# save_last=False, save_last=False,
# mode="min", mode="min",
# ) )
# custom_ckpt = CustomCheckpointIO() custom_ckpt = CustomCheckpointIO()
# trainer = pl.Trainer( trainer = pl.Trainer(
# accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
# devices=1, devices=1,
# max_epochs=config.train_config.max_epochs, max_epochs=config.train_config.max_epochs,
# val_check_interval=config.train_config.val_check_interval, val_check_interval=config.train_config.val_check_interval,
# check_val_every_n_epoch=config.train_config.check_val_every_n_epoch, check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
# gradient_clip_val=config.train_config.gradient_clip_val, gradient_clip_val=config.train_config.gradient_clip_val,
# precision=16, # we'll use mixed precision precision=16, # we'll use mixed precision
# plugins=custom_ckpt, plugins=custom_ckpt,
# num_sanity_val_steps=0, num_sanity_val_steps=0,
# logger=wandb_logger, logger=wandb_logger,
# callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback], callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback],
# ) )
# trainer.fit(model_module) trainer.fit(model_module)
if __name__ == "__main__": if __name__ == "__main__":

2
utils/donut_dataset_stream.py
View File

@ -47,7 +47,7 @@ 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 = dataset self.dataset = dataset[self.split]
self.dataset_length = len(self.dataset) self.dataset_length = len(self.dataset)
self.gt_token_sequences = [] self.gt_token_sequences = []