onmt/train_single.py

#!/usr/bin/env python
"""Training on a single process."""
import torch
import sys

from onmt.utils.logging import init_logger, logger
from onmt.utils.parse import ArgumentParser
from onmt.constants import CorpusTask
from onmt.transforms import (
    make_transforms,
    save_transforms,
    get_specials,
    get_transforms_cls,
)
from onmt.inputters import build_vocab, IterOnDevice
from onmt.inputters.inputter import dict_to_vocabs, vocabs_to_dict
from onmt.inputters.dynamic_iterator import build_dynamic_dataset_iter
from onmt.inputters.text_corpus import save_transformed_sample
from onmt.model_builder import build_model
from onmt.models.model_saver import load_checkpoint
from onmt.utils.optimizers import Optimizer
from onmt.utils.misc import set_random_seed
from onmt.trainer import build_trainer
from onmt.models import build_model_saver
from onmt.modules.embeddings import prepare_pretrained_embeddings


def prepare_transforms_vocabs(opt, transforms_cls):
    """Prepare or dump transforms before training."""
    # if transform + options set in 'valid' we need to copy in main
    # transform / options for scoring considered as inference
    validset_transforms = opt.data.get("valid", {}).get("transforms", None)
    if validset_transforms:
        opt.transforms = validset_transforms
        if opt.data.get("valid", {}).get("tgt_prefix", None):
            opt.tgt_prefix = opt.data.get("valid", {}).get("tgt_prefix", None)
        if opt.data.get("valid", {}).get("src_prefix", None):
            opt.src_prefix = opt.data.get("valid", {}).get("src_prefix", None)
        if opt.data.get("valid", {}).get("tgt_suffix", None):
            opt.tgt_suffix = opt.data.get("valid", {}).get("tgt_suffix", None)
        if opt.data.get("valid", {}).get("src_suffix", None):
            opt.src_suffix = opt.data.get("valid", {}).get("src_suffix", None)
    specials = get_specials(opt, transforms_cls)

    vocabs = build_vocab(opt, specials)

    # maybe prepare pretrained embeddings, if any
    prepare_pretrained_embeddings(opt, vocabs)

    if opt.dump_transforms or opt.n_sample != 0:
        transforms = make_transforms(opt, transforms_cls, vocabs)
    if opt.dump_transforms:
        save_transforms(transforms, opt.save_data, overwrite=opt.overwrite)
    if opt.n_sample != 0:
        logger.warning(
            "`-n_sample` != 0: Training will not be started. "
            f"Stop after saving {opt.n_sample} samples/corpus."
        )
        save_transformed_sample(opt, transforms, n_sample=opt.n_sample)
        logger.info("Sample saved, please check it before restart training.")
        sys.exit()
    logger.info(
        "The first 10 tokens of the vocabs are:"
        f"{vocabs_to_dict(vocabs)['src'][0:10]}"
    )
    logger.info(f"The decoder start token is: {opt.decoder_start_token}")
    return vocabs


def _init_train(opt):
    """Common initilization stuff for all training process.
    We need to build or rebuild the vocab in 3 cases:
    - training from scratch (train_from is false)
    - resume training but transforms have changed
    - resume training but vocab file has been modified
    """
    ArgumentParser.validate_prepare_opts(opt)
    transforms_cls = get_transforms_cls(opt._all_transform)
    if opt.train_from:
        # Load checkpoint if we resume from a previous training.
        checkpoint = load_checkpoint(ckpt_path=opt.train_from)
        vocabs = dict_to_vocabs(checkpoint["vocab"])
        if (
            hasattr(checkpoint["opt"], "_all_transform")
            and len(
                opt._all_transform.symmetric_difference(
                    checkpoint["opt"]._all_transform
                )
            )
            != 0
        ):
            _msg = "configured transforms is different from checkpoint:"
            new_transf = opt._all_transform.difference(checkpoint["opt"]._all_transform)
            old_transf = checkpoint["opt"]._all_transform.difference(opt._all_transform)
            if len(new_transf) != 0:
                _msg += f" +{new_transf}"
            if len(old_transf) != 0:
                _msg += f" -{old_transf}."
            logger.warning(_msg)
            vocabs = prepare_transforms_vocabs(opt, transforms_cls)
        if opt.update_vocab:
            logger.info("Updating checkpoint vocabulary with new vocabulary")
            vocabs = prepare_transforms_vocabs(opt, transforms_cls)
    else:
        checkpoint = None
        vocabs = prepare_transforms_vocabs(opt, transforms_cls)

    return checkpoint, vocabs, transforms_cls


def configure_process(opt, device_id):
    if device_id >= 0:
        torch.cuda.set_device(device_id)
    set_random_seed(opt.seed, device_id >= 0)


def _get_model_opts(opt, checkpoint=None):
    """Get `model_opt` to build model, may load from `checkpoint` if any."""
    if checkpoint is not None:
        model_opt = ArgumentParser.ckpt_model_opts(checkpoint["opt"])
        if opt.override_opts:
            logger.info("Over-ride model option set to true - use with care")
            args = list(opt.__dict__.keys())
            model_args = list(model_opt.__dict__.keys())
            for arg in args:
                if arg in model_args and getattr(opt, arg) != getattr(model_opt, arg):
                    logger.info(
                        "Option: %s , value: %s overriding model: %s"
                        % (arg, getattr(opt, arg), getattr(model_opt, arg))
                    )
            model_opt = opt
        else:
            model_opt = ArgumentParser.ckpt_model_opts(checkpoint["opt"])
            ArgumentParser.update_model_opts(model_opt)
            ArgumentParser.validate_model_opts(model_opt)
            if opt.tensorboard_log_dir == model_opt.tensorboard_log_dir and hasattr(
                model_opt, "tensorboard_log_dir_dated"
            ):
                # ensure tensorboard output is written in the directory
                # of previous checkpoints
                opt.tensorboard_log_dir_dated = (
                    model_opt.tensorboard_log_dir_dated
                )  # noqa: E501
            # Override checkpoint's update_embeddings as it defaults to false
            model_opt.update_vocab = opt.update_vocab
            # Override checkpoint's freezing settings as it defaults to false
            model_opt.freeze_encoder = opt.freeze_encoder
            model_opt.freeze_decoder = opt.freeze_decoder
    else:
        model_opt = opt
    return model_opt


def main(opt, device_id):
    """Start training on `device_id`."""
    # NOTE: It's important that ``opt`` has been validated and updated
    # at this point.

    configure_process(opt, device_id)
    init_logger(opt.log_file)
    checkpoint, vocabs, transforms_cls = _init_train(opt)
    model_opt = _get_model_opts(opt, checkpoint=checkpoint)

    # Build model.
    model = build_model(model_opt, opt, vocabs, checkpoint, device_id)

    model.count_parameters(log=logger.info)
    trainable = {
        "torch.float32": 0,
        "torch.float16": 0,
        "torch.uint8": 0,
        "torch.int8": 0,
    }
    non_trainable = {
        "torch.float32": 0,
        "torch.float16": 0,
        "torch.uint8": 0,
        "torch.int8": 0,
    }
    for n, p in model.named_parameters():
        if p.requires_grad:
            trainable[str(p.dtype)] += p.numel()
        else:
            non_trainable[str(p.dtype)] += p.numel()
    logger.info("Trainable parameters = %s" % str(trainable))
    logger.info("Non trainable parameters = %s" % str(non_trainable))
    logger.info(" * src vocab size = %d" % len(vocabs["src"]))
    logger.info(" * tgt vocab size = %d" % len(vocabs["tgt"]))
    if "src_feats" in vocabs:
        for i, feat_vocab in enumerate(vocabs["src_feats"]):
            logger.info(f"* src_feat {i} vocab size = {len(feat_vocab)}")

    # Build optimizer.
    optim = Optimizer.from_opt(model, opt, checkpoint=checkpoint)

    del checkpoint

    # Build model saver
    model_saver = build_model_saver(model_opt, opt, model, vocabs, optim, device_id)

    trainer = build_trainer(
        opt, device_id, model, vocabs, optim, model_saver=model_saver
    )

    offset = max(0, device_id) if opt.parallel_mode == "data_parallel" else 0
    stride = max(1, len(opt.gpu_ranks)) if opt.parallel_mode == "data_parallel" else 1

    _train_iter = build_dynamic_dataset_iter(
        opt,
        transforms_cls,
        vocabs,
        task=CorpusTask.TRAIN,
        copy=opt.copy_attn,
        stride=stride,
        offset=offset,
    )
    train_iter = IterOnDevice(_train_iter, device_id)

    valid_iter = build_dynamic_dataset_iter(
        opt, transforms_cls, vocabs, task=CorpusTask.VALID, copy=opt.copy_attn
    )

    if valid_iter is not None:
        valid_iter = IterOnDevice(valid_iter, device_id)

    if len(opt.gpu_ranks):
        logger.info("Starting training on GPU: %s" % opt.gpu_ranks)
    else:
        logger.info("Starting training on CPU, could be very slow")
    train_steps = opt.train_steps
    if opt.single_pass and train_steps > 0:
        logger.warning("Option single_pass is enabled, ignoring train_steps.")
        train_steps = 0

    trainer.train(
        train_iter,
        train_steps,
        save_checkpoint_steps=opt.save_checkpoint_steps,
        valid_iter=valid_iter,
        valid_steps=opt.valid_steps,
    )

    if trainer.report_manager.tensorboard_writer is not None:
        trainer.report_manager.tensorboard_writer.close()