testing
This commit is contained in:
parent
93a231a477
commit
dd5febad65
108
train_stream.py
108
train_stream.py
@ -34,68 +34,72 @@ def main(config, hug_token):
|
||||
|
||||
added_tokens = []
|
||||
|
||||
train_dataset = DonutDataset(
|
||||
config.dataset_path,
|
||||
processor=processor,
|
||||
model=model,
|
||||
max_length=config.max_length,
|
||||
split="train",
|
||||
task_start_token="<s_cord-v2>",
|
||||
prompt_end_token="<s_cord-v2>",
|
||||
added_tokens=added_tokens,
|
||||
sort_json_key=False, # cord dataset is preprocessed, so no need for this
|
||||
)
|
||||
dataset = load_dataset(config.dataset_path)
|
||||
dataset.train_test_split(test_size=0.1)
|
||||
print(dataset)
|
||||
|
||||
val_dataset = DonutDataset(
|
||||
config.dataset_path,
|
||||
processor=processor,
|
||||
model=model,
|
||||
max_length=config.max_length,
|
||||
split="validation",
|
||||
task_start_token="<s_cord-v2>",
|
||||
prompt_end_token="<s_cord-v2>",
|
||||
added_tokens=added_tokens,
|
||||
sort_json_key=False, # cord dataset is preprocessed, so no need for this
|
||||
)
|
||||
# train_dataset = DonutDataset(
|
||||
# dataset,
|
||||
# processor=processor,
|
||||
# model=model,
|
||||
# max_length=config.max_length,
|
||||
# split="train",
|
||||
# task_start_token="<s_cord-v2>",
|
||||
# prompt_end_token="<s_cord-v2>",
|
||||
# added_tokens=added_tokens,
|
||||
# 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.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
|
||||
# val_dataset = DonutDataset(
|
||||
# dataset,
|
||||
# processor=processor,
|
||||
# model=model,
|
||||
# max_length=config.max_length,
|
||||
# split="validation",
|
||||
# task_start_token="<s_cord-v2>",
|
||||
# prompt_end_token="<s_cord-v2>",
|
||||
# added_tokens=added_tokens,
|
||||
# sort_json_key=False, # cord dataset is preprocessed, so no need for this
|
||||
# )
|
||||
|
||||
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)
|
||||
# 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]
|
||||
|
||||
login(hug_token, True)
|
||||
# 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)
|
||||
|
||||
model_module = DonutModelPLModule(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader)
|
||||
# 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)
|
||||
|
||||
wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)
|
||||
# wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)
|
||||
|
||||
checkpoint_callback = ModelCheckpoint(
|
||||
monitor="val_metric",
|
||||
dirpath=config.checkpoint_path,
|
||||
filename="artifacts",
|
||||
save_top_k=1,
|
||||
save_last=False,
|
||||
mode="min",
|
||||
)
|
||||
# checkpoint_callback = ModelCheckpoint(
|
||||
# monitor="val_metric",
|
||||
# dirpath=config.checkpoint_path,
|
||||
# filename="artifacts",
|
||||
# save_top_k=1,
|
||||
# save_last=False,
|
||||
# mode="min",
|
||||
# )
|
||||
|
||||
custom_ckpt = CustomCheckpointIO()
|
||||
# custom_ckpt = CustomCheckpointIO()
|
||||
|
||||
trainer = pl.Trainer(
|
||||
accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
|
||||
devices=1,
|
||||
max_epochs=config.train_config.max_epochs,
|
||||
val_check_interval=config.train_config.val_check_interval,
|
||||
check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
|
||||
gradient_clip_val=config.train_config.gradient_clip_val,
|
||||
precision=16, # we'll use mixed precision
|
||||
plugins=custom_ckpt,
|
||||
num_sanity_val_steps=0,
|
||||
logger=wandb_logger,
|
||||
callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback],
|
||||
)
|
||||
# trainer = pl.Trainer(
|
||||
# accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
|
||||
# devices=1,
|
||||
# max_epochs=config.train_config.max_epochs,
|
||||
# val_check_interval=config.train_config.val_check_interval,
|
||||
# check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
|
||||
# gradient_clip_val=config.train_config.gradient_clip_val,
|
||||
# precision=16, # we'll use mixed precision
|
||||
# plugins=custom_ckpt,
|
||||
# num_sanity_val_steps=0,
|
||||
# logger=wandb_logger,
|
||||
# callbacks=[PushToHubCallback(output_model_path=config.output_model_path), checkpoint_callback],
|
||||
# )
|
||||
|
||||
trainer.fit(model_module)
|
||||
# trainer.fit(model_module)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
@ -24,7 +24,7 @@ class DonutDataset(Dataset):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_name_or_path: str,
|
||||
dataset: Dataset,
|
||||
max_length: int,
|
||||
processor: DonutProcessor,
|
||||
model: VisionEncoderDecoderModel,
|
||||
@ -47,8 +47,7 @@ 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 = dataset
|
||||
self.dataset_length = len(self.dataset)
|
||||
|
||||
self.gt_token_sequences = []
|
||||
|
Loading…
Reference in New Issue
Block a user