DimaKoshman's picture
trained model for a bit
028951c
raw
history blame
3.53 kB
import os
import pandas as pd
import pytorch_lightning as pl
import transformers
import wandb
from config import CONFIG
from data import (
get_annotation_ground_truth_str_from_image_index,
load_train_image_ids,
build_dataloader,
Split,
Batch,
)
from metrics import benetech_score_string_prediction
from model import generate_token_strings, LightningModule
from utils import set_tokenizers_parallelism, set_torch_device_order_pci_bus
class MetricsCallback(pl.callbacks.Callback):
def on_validation_batch_start(
self, trainer, pl_module, batch: Batch, batch_idx, dataloader_idx=0
):
predicted_strings = generate_token_strings(pl_module.model, images=batch.images)
for expected_data_index, predicted_string in zip(
batch.data_indices, predicted_strings, strict=True
):
benetech_score = benetech_score_string_prediction(
expected_data_index=expected_data_index,
predicted_string=predicted_string,
)
wandb.log(dict(benetech_score=benetech_score))
ground_truth_strings = [
get_annotation_ground_truth_str_from_image_index(i)
for i in batch.data_indices
]
string_ids = [load_train_image_ids()[i] for i in batch.data_indices]
strings_dataframe = pd.DataFrame(
dict(
string_ids=string_ids,
ground_truth=ground_truth_strings,
predicted=predicted_strings,
)
)
wandb.log(dict(strings=wandb.Table(dataframe=strings_dataframe)))
class TransformersPreTrainedModelsCheckpointIO(pl.plugins.CheckpointIO):
def __init__(
self, pretrained_models: list[transformers.modeling_utils.PreTrainedModel]
):
super().__init__()
self.pretrained_models = pretrained_models
def save_checkpoint(self, checkpoint, path, storage_options=None):
for pretrained_model in self.pretrained_models:
pretrained_model.save_pretrained(path)
def load_checkpoint(self, path, storage_options=None):
self.pretrained_models = [
pm.from_pretrained(path) for pm in self.pretrained_models
]
def remove_checkpoint(self, path):
os.remove(path)
def train():
set_tokenizers_parallelism(False)
set_torch_device_order_pci_bus()
pl_module = LightningModule(CONFIG)
model_checkpoint = pl.callbacks.ModelCheckpoint(
dirpath=CONFIG.training_directory,
monitor="val_loss",
save_top_k=CONFIG.save_top_k_checkpoints,
)
metrics_callback = MetricsCallback()
logger = pl.loggers.WandbLogger(
project=CONFIG.wandb_project_name, save_dir=CONFIG.training_directory
)
plugin = TransformersPreTrainedModelsCheckpointIO(
[pl_module.model.processor, pl_module.model.encoder_decoder]
)
trainer = pl.Trainer(
accelerator=CONFIG.accelerator,
devices=CONFIG.devices,
plugins=[plugin],
callbacks=[model_checkpoint, metrics_callback],
logger=logger,
limit_train_batches=CONFIG.limit_train_batches,
limit_val_batches=CONFIG.limit_val_batches,
)
trainer.fit(
model=pl_module,
train_dataloaders=build_dataloader(
Split.train, pl_module.model.batch_collate_function
),
val_dataloaders=build_dataloader(
Split.val, pl_module.model.batch_collate_function
),
)
if __name__ == "__main__":
train()