pylaia-htr-train-ctc

#!/usr/bin/env python
from __future__ import absolute_import

import argparse
import multiprocessing
import os
import random

import numpy
import torch
from torch.optim import RMSprop

import laia.common.logging as log
import laia.data.transforms as transforms
from laia.common.arguments import add_argument, args, add_defaults
from laia.common.arguments_types import NumberInClosedRange, str2bool
from laia.common.loader import ModelLoader, StateCheckpointLoader
from laia.common.random import manual_seed
from laia.common.saver import (
    CheckpointSaver,
    RollingSaver,
    ModelCheckpointSaver,
    StateCheckpointSaver,
)
from laia.conditions import Lowest, MultipleOf, GEqThan, ConsecutiveNonDecreasing
from laia.data import ImageDataLoader, TextImageFromTextTableDataset, FixedSizeSampler
from laia.engine import Trainer, Evaluator
from laia.engine.engine import EPOCH_END, EPOCH_START
from laia.engine.feeders import ImageFeeder, ItemFeeder
from laia.experiments.htr_experiment import HTRExperiment
from laia.hooks import Hook, HookList, action, Action
from laia.losses.ctc_loss import (
    CTCLossImpl,
    get_default_add_logsoftmax,
    set_default_add_logsoftmax,
    set_default_implementation,
)
from laia.utils import SymbolsTable


def worker_init_fn(_):
    # We need to reset the Numpy and Python PRNG, or we will get the
    # same numbers in each epoch (when the workers are re-generated)
    random.seed(torch.initial_seed() % 2 ** 31)
    numpy.random.seed(torch.initial_seed() % 2 ** 31)


if __name__ == "__main__":
    add_defaults(
        "batch_size",
        "learning_rate",
        "momentum",
        "gpu",
        "max_epochs",
        "seed",
        "show_progress_bar",
        "train_path",
        "train_samples_per_epoch",
        "valid_samples_per_epoch",
        "iterations_per_update",
        "save_checkpoint_interval",
        "num_rolling_checkpoints",
        "use_distortions",
    )
    add_argument(
        "syms",
        type=argparse.FileType("r"),
        help="Symbols table mapping from strings to integers.",
    )
    add_argument(
        "img_dirs", type=str, nargs="+", help="Directory containing word images."
    )
    add_argument(
        "tr_txt_table",
        type=argparse.FileType("r"),
        help="Character transcriptions of each training image.",
    )
    add_argument(
        "va_txt_table",
        type=argparse.FileType("r"),
        help="Character transcriptions of each validation image.",
    )
    add_argument(
        "--delimiters",
        type=str,
        nargs="+",
        default=["<space>"],
        help="Sequence of characters representing the word delimiters.",
    )
    add_argument(
        "--max_nondecreasing_epochs",
        type=NumberInClosedRange(int, vmin=0),
        help="Stop the training once there has been this number "
        "consecutive epochs without a new lowest validation CER.",
    )
    add_argument(
        "--model_filename", type=str, default="model", help="File name of the model."
    )
    add_argument(
        "--checkpoint",
        type=str,
        default="ckpt.lowest-valid-cer*",
        help="Suffix of the checkpoint to use, can be a glob pattern.",
    )
    add_argument(
        "--use_baidu_ctc",
        type=str2bool,
        nargs="?",
        const=True,
        default=False,
        help="If true, use Baidu's implementation of the CTC loss.",
    )
    add_argument(
        "--add_logsoftmax_to_loss",
        type=str2bool,
        nargs="?",
        const=True,
        default=get_default_add_logsoftmax(),
        help="If true, add a logsoftmax operation before the CTC loss to normalize the activations.",
    )
    args = args()

    manual_seed(args.seed)
    syms = SymbolsTable(args.syms)
    device = torch.device("cuda:{}".format(args.gpu - 1) if args.gpu else "cpu")

    # Set the default options for the CTCLoss.
    set_default_add_logsoftmax(args.add_logsoftmax_to_loss)
    if args.use_baidu_ctc:
        set_default_implementation(CTCLossImpl.BAIDU)

    model = ModelLoader(
        args.train_path, filename=args.model_filename, device=device
    ).load()
    if model is None:
        log.error('Could not find the model. Have you run "pylaia-htr-create-model"?')
        exit(1)
    model = model.to(device)

    default_img_transform = transforms.Compose(
        [
            transforms.vision.Convert("L"),
            transforms.vision.Invert(),
            transforms.vision.ToTensor(),
        ]
    )
    if args.use_distortions:
        tr_img_transform = transforms.Compose(
            [
                transforms.vision.Convert("L"),
                transforms.vision.Invert(),
                transforms.vision.RandomBetaAffine(),
                transforms.vision.ToTensor(),
            ]
        )
    else:
        tr_img_transform = default_img_transform
    log.info("Training data transforms:\n{}", str(tr_img_transform))

    tr_dataset = TextImageFromTextTableDataset(
        args.tr_txt_table,
        args.img_dirs,
        img_transform=tr_img_transform,
        txt_transform=transforms.text.ToTensor(syms),
    )
    tr_dataset_loader = ImageDataLoader(
        dataset=tr_dataset,
        image_channels=1,
        batch_size=args.batch_size,
        num_workers=multiprocessing.cpu_count(),
        shuffle=not bool(args.train_samples_per_epoch),
        sampler=FixedSizeSampler(tr_dataset, args.train_samples_per_epoch)
        if args.train_samples_per_epoch
        else None,
        worker_init_fn=worker_init_fn,
    )
    trainer = Trainer(
        model=model,
        criterion=None,  # Set automatically by HTRExperiment
        optimizer=RMSprop(
            model.parameters(), lr=args.learning_rate, momentum=args.momentum
        ),
        data_loader=tr_dataset_loader,
        batch_input_fn=ImageFeeder(device=device, parent_feeder=ItemFeeder("img")),
        batch_target_fn=ItemFeeder("txt"),
        batch_id_fn=ItemFeeder("id"),  # Print image ids on exception
        progress_bar="Train" if args.show_progress_bar else None,
        iterations_per_update=args.iterations_per_update,
    )

    va_dataset = TextImageFromTextTableDataset(
        args.va_txt_table,
        args.img_dirs,
        img_transform=default_img_transform,
        txt_transform=transforms.text.ToTensor(syms),
    )
    va_dataset_loader = ImageDataLoader(
        dataset=va_dataset,
        image_channels=1,
        batch_size=args.batch_size,
        num_workers=multiprocessing.cpu_count(),
        sampler=FixedSizeSampler(va_dataset, args.valid_samples_per_epoch)
        if args.valid_samples_per_epoch
        else None,
    )
    evaluator = Evaluator(
        model=model,
        data_loader=va_dataset_loader,
        batch_input_fn=ImageFeeder(device=device, parent_feeder=ItemFeeder("img")),
        batch_target_fn=ItemFeeder("txt"),
        batch_id_fn=ItemFeeder("id"),
        progress_bar="Valid" if args.show_progress_bar else None,
    )

    experiment = HTRExperiment(
        trainer, evaluator, word_delimiters=[syms[sym] for sym in args.delimiters]
    )

    def ckpt_saver(filename, obj):
        return RollingSaver(
            StateCheckpointSaver(
                CheckpointSaver(os.path.join(args.train_path, filename)),
                obj,
                device=device,
            ),
            keep=args.num_rolling_checkpoints,
        )

    saver_best_cer = ckpt_saver("experiment.ckpt.lowest-valid-cer", experiment)
    saver_best_wer = ckpt_saver("experiment.ckpt.lowest-valid-wer", experiment)

    @action
    def save(saver, epoch):
        saver.save(suffix=epoch)

    # Set hooks
    trainer.add_hook(
        EPOCH_END,
        HookList(
            # Save on best CER
            Hook(Lowest(experiment.valid_cer()), Action(save, saver=saver_best_cer)),
            # Save on best WER
            Hook(Lowest(experiment.valid_wer()), Action(save, saver=saver_best_wer)),
        ),
    )
    if args.save_checkpoint_interval:
        # Save every `save_checkpoint_interval` epochs
        log.get_logger("laia.hooks.conditions.multiple_of").setLevel(log.WARNING)
        trainer.add_hook(
            EPOCH_END,
            Hook(
                MultipleOf(trainer.epochs, args.save_checkpoint_interval),
                Action(save, saver=ckpt_saver("experiment.ckpt", experiment)),
            ),
        )
    if args.max_nondecreasing_epochs:
        # Stop when the validation CER hasn't improved in
        # `max_nondecreasing_epochs` consecutive epochs
        trainer.add_hook(
            EPOCH_END,
            Hook(
                ConsecutiveNonDecreasing(
                    experiment.valid_cer(), args.max_nondecreasing_epochs
                ),
                trainer.stop,
            ),
        )
    if args.max_epochs:
        # Stop when `max_epochs` has been reached
        trainer.add_hook(
            EPOCH_START, Hook(GEqThan(trainer.epochs, args.max_epochs), trainer.stop)
        )

    # Continue from the given checkpoint, if possible
    StateCheckpointLoader(experiment, device=device).load_by(
        os.path.join(args.train_path, "experiment.{}".format(args.checkpoint))
    )

    experiment.run()

    # Experiment finished. Save the model separately
    ModelCheckpointSaver(
        CheckpointSaver(os.path.join(args.train_path, "model.ckpt")), model
    ).save(suffix="last")