run_training.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.

from datetime import datetime
import logging
import math
import os
import shutil
import tempfile
from pathlib import Path
import random
import string

import tensorflow as tf
import wandb
import tensorflow.keras.backend as K
from tensorflow.python.ipu import optimizers

import utils
import xpu
from argparser import parse_args
from custom_callbacks import CheckpointCallback
from data_utils.input_spec import create_inputs_from_features
from data_utils.load_dataset import load_raw_dataset
from data_utils.preprocess_dataset import preprocess_dataset
from inference import run_inference
from utils import ThroughputCallback, get_optimizer, str_dtype_to_tf_dtype, convert_loss_and_metric_reductions_to_fp32, options_validator
from model.utils import (create_model, get_loss_functions, get_metrics, get_tf_dataset)
from pipeline.pipeline_stage_assignment import pipeline_model
from pipeline.pipeline_stage_names import (PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES)
# Config Options
logging.basicConfig(level=logging.INFO)
cfg = parse_args()
options_validator(cfg)


def cosine_lr(epoch, lr):
    completed_fraction = epoch / cfg.model.epochs
    cosine = tf.cos(tf.constant(math.pi, dtype=tf.float32) * completed_fraction)
    cosine_decayed = 0.5 * (1.0 + cosine)
    lr = max(cosine_decayed * cfg.model.lr, cfg.model.min_lr)
    if cfg.model.lr_warmup_epochs > 0 and epoch < cfg.model.lr_warmup_epochs:
        init_prop = cfg.model.lr_init_prop
        w_ep = cfg.model.lr_warmup_epochs
        lr_scale = (init_prop * (w_ep - epoch) / w_ep) + (epoch / w_ep)
        lr *= lr_scale
    if cfg.wandb:
        wandb.log({"lr": lr})
    return lr


def linear_lr(epoch, lr):
    warmup = cfg.model.lr_warmup_epochs
    if warmup > 0 and epoch < warmup:
        lr = epoch * cfg.model.lr / warmup
    else:
        p = (epoch - warmup) / (cfg.model.epochs - warmup)
        lr = p * cfg.model.min_lr + (1 - p) * cfg.model.lr
    if cfg.wandb:
        wandb.log({"lr": lr})
    return lr


def main():
    if cfg.wandb:
        wandb.init(entity=cfg.wandb_entity, project=cfg.wandb_project, config=cfg.as_dict())
    tf.keras.mixed_precision.set_global_policy(cfg.model.dtype)
    # ------------ CHECKPOINTS ---------------
    # Create a directory for the checkpoints based on model name and time
    # Want to check directory exisits before lengthy data processing
    if cfg.save_checkpoints_locally:
        if not Path(cfg.checkpoint_dir).is_dir():
            os.mkdir(Path(cfg.checkpoint_dir))
        assert Path(cfg.checkpoint_dir).is_dir()
        date_string = 'model-' + datetime.now().strftime("%d-%m-%Y_%H-%M-%S") + "-" + \
                        ''.join(random.choice(string.ascii_lowercase) for _ in range(10))
        model_dir = os.path.join(cfg.checkpoint_dir, date_string)
    else:
        # If no checkpoint specified temp path is given
        model_dir = tempfile.mkdtemp()
    logging.info(f"Model weights will be saved to {model_dir}")

    logging.info(f"Dataset: {cfg.dataset.dataset_name}")
    graph_data = load_raw_dataset(cfg.dataset.dataset_name, Path(cfg.dataset.cache_path), cfg)

    graph_data = preprocess_dataset(dataset=graph_data, options=cfg)

    graph_data.get_statistics(do_check_for_nans=cfg.debug.check_data_for_nans)
    if cfg.debug.calc_conformer_stats:
        graph_data.get_conformer_statistics()

    input_spec = create_inputs_from_features(dataset=graph_data, cfg=cfg, fold="train")

    batch_generator, _ = get_tf_dataset(preprocessed_dataset=graph_data,
                                        split_name="train",
                                        shuffle=True,
                                        options=cfg,
                                        input_spec=input_spec)

    if cfg.model.target_total_batch_size:
        if cfg.model.micro_batch_size and (not cfg.ipu_opts.gradient_accumulation_factor):
            fixed_batch_size = float(cfg.model.micro_batch_size) * \
                               float(cfg.ipu_opts.replicas) * \
                               batch_generator.stats['avg_pack']['graphs']
            new_GA = max(int(round(cfg.model.target_total_batch_size / fixed_batch_size)), 1)
            logging.info(f"Setting gradient_accumulation_factor to {new_GA}")
            cfg.ipu_opts.gradient_accumulation_factor = new_GA
        else:
            fixed_batch_size = float(cfg.ipu_opts.gradient_accumulation_factor) * \
                               float(cfg.ipu_opts.replicas) * \
                               batch_generator.stats['avg_pack']['graphs']
            new_micro_batch_size = int(round(cfg.model.target_total_batch_size / fixed_batch_size))
            new_micro_batch_size = max(new_micro_batch_size, 1)
            logging.info(f"Changing micro_batch_size from {cfg.model.micro_batch_size} to {new_micro_batch_size}")
            cfg.model.micro_batch_size = new_micro_batch_size
            batch_generator.change_micro_batch_size(cfg.model.micro_batch_size)

        batch_generator.get_averaged_global_batch_size(cfg.model.micro_batch_size,
                                                       cfg.ipu_opts.gradient_accumulation_factor, cfg.ipu_opts.replicas)
        logging.info(f"Dataset stats: {batch_generator.stats}")

    if cfg.ipu_opts.num_pipeline_stages > 1:
        round_to = float(2 * cfg.ipu_opts.num_pipeline_stages)
        new_GA = int(max(round(cfg.ipu_opts.gradient_accumulation_factor / round_to), 1) * round_to)
        if new_GA != cfg.ipu_opts.gradient_accumulation_factor:
            logging.info(f"Rounding gradient_accumulation_factor to a multiple of {round_to}: {new_GA}")
            cfg.ipu_opts.gradient_accumulation_factor = new_GA
            batch_generator.get_averaged_global_batch_size(
                cfg.model.micro_batch_size, cfg.ipu_opts.gradient_accumulation_factor, cfg.ipu_opts.replicas)
            logging.info(f"Dataset stats: {batch_generator.stats}")

    if cfg.wandb:
        try:
            wandb.log({'dataset_stats': batch_generator.stats})
        except AttributeError:
            logging.info(f"Batch generator {batch_generator} has no stats object.")

    tf_dataset = batch_generator.get_tf_dataset()
    steps_per_epoch = batch_generator.batches_per_epoch
    steps_per_execution_per_replica = steps_per_epoch // cfg.ipu_opts.replicas
    steps_per_execution_per_replica = cfg.ipu_opts.gradient_accumulation_factor * (
        steps_per_execution_per_replica // cfg.ipu_opts.gradient_accumulation_factor)
    new_steps_per_epoch = steps_per_execution_per_replica * cfg.ipu_opts.replicas
    if new_steps_per_epoch != steps_per_epoch:
        logging.warning("Steps per epoch has been truncated from"
                        f" {steps_per_epoch} to {new_steps_per_epoch}"
                        " in order for it to be divisible by steps per execution.")
    steps_per_epoch = new_steps_per_epoch
    logging.info(f"Steps per epoch: {steps_per_epoch}")
    logging.info(f"Steps per execution per replica: {steps_per_execution_per_replica}")

    optimizer_options = dict(
        name=cfg.model.opt.lower(),
        learning_rate=cfg.model.lr,
        l2_regularization=cfg.model.l2_regularization,
        dtype=str_dtype_to_tf_dtype(cfg.model.dtype),
        m_dtype=str_dtype_to_tf_dtype(cfg.model.adam_m_dtype),
        v_dtype=str_dtype_to_tf_dtype(cfg.model.adam_v_dtype),
        clip_value=cfg.model.grad_clip_value,
        loss_scale=cfg.model.loss_scaling,
        gradient_accumulation_factor=cfg.ipu_opts.gradient_accumulation_factor,
        replicas=cfg.ipu_opts.replicas,
        outline_apply_gradients=not cfg.ipu_opts.offload_optimizer_state,  # bug where outlining causes issues
    )

    # ------------ TRAINING LOOP ---------------
    if cfg.do_training:
        num_pipeline_stages_training = cfg.ipu_opts.num_pipeline_stages
        strategy_training = xpu.configure_and_get_strategy(num_replicas=cfg.ipu_opts.replicas,
                                                           num_ipus_per_replica=num_pipeline_stages_training,
                                                           cfg=cfg)
        with strategy_training.scope():
            model = create_model(batch_generator, graph_data, cfg, input_spec=input_spec)
            model.summary()
            utils.print_trainable_variables(model, log_wandb=cfg.wandb)
            if num_pipeline_stages_training > 1:
                pipeline_model(model=model,
                               config=cfg,
                               pipeline_names=PIPELINE_NAMES,
                               pipeline_allocate_previous=PIPELINE_ALLOCATE_PREVIOUS,
                               num_pipeline_stages=num_pipeline_stages_training,
                               matmul_partials_type='half')

            losses, loss_weights = get_loss_functions(graph_data, cfg)
            metrics = get_metrics(graph_data.denormalize, cfg)

            callbacks = [
                ThroughputCallback(
                    # the throughput depends on the COMPUTE batch size, not the TOTAL batch size
                    samples_per_epoch=batch_generator.n_graphs_per_epoch,
                    log_wandb=cfg.wandb)
            ]
            if cfg.model.learning_rate_schedule == 'cosine':
                callbacks.append(tf.keras.callbacks.LearningRateScheduler(cosine_lr))
            elif cfg.model.learning_rate_schedule == 'linear':
                callbacks.append(tf.keras.callbacks.LearningRateScheduler(linear_lr))
            if cfg.execution_profile:
                callbacks.append(tf.keras.callbacks.TensorBoard(profile_batch=[2], log_dir='logs'))
            if cfg.wandb:
                callbacks.append(wandb.keras.WandbCallback())

            logging.info("Running training...")
            logging.info(f"Saving weights to {model_dir}")
            model_path = os.path.join(model_dir, 'model-{epoch:05d}')
            callbacks.append(
                CheckpointCallback(use_wandb=cfg.wandb,
                                   upload_to_wandb=cfg.upload_final_ckpt,
                                   save_checkpoints_locally=cfg.save_checkpoints_locally,
                                   filepath=model_path,
                                   verbose=1,
                                   save_best_only=False,
                                   save_weights_only=True,
                                   period=cfg.checkpoint_every_n_epochs,
                                   total_epochs=cfg.model.epochs,
                                   model_name='model'))

            model.compile(optimizer=get_optimizer(**optimizer_options),
                          loss=losses,
                          loss_weights=loss_weights,
                          weighted_metrics=metrics,
                          steps_per_execution=steps_per_execution_per_replica)
            if cfg.model.dtype == 'float16':
                # the loss reduction is set by backend.floatx by default
                # must be forced to reduce in float32 to avoid overflow
                convert_loss_and_metric_reductions_to_fp32(model)

            # if the total batch size exceeds the compute batch size
            if xpu.IS_IPU:
                model.set_gradient_accumulation_options(
                    gradient_accumulation_steps_per_replica=cfg.ipu_opts.gradient_accumulation_factor,
                    gradient_accumulation_reduction_method=optimizers.GradientAccumulationReductionMethod.RUNNING_MEAN,
                    dtype=str_dtype_to_tf_dtype(cfg.ipu_opts.gradient_accumulation_dtype or cfg.model.dtype),
                    offload_weight_update_variables=cfg.ipu_opts.offload_optimizer_state,
                    replicated_optimizer_state_sharding=cfg.ipu_opts.RTS)
            model.fit(tf_dataset, steps_per_epoch=steps_per_epoch, epochs=cfg.model.epochs, callbacks=callbacks)

    if cfg.checkpoint_path:
        checkpoint_paths = {-1: cfg.checkpoint_path}
    else:
        checkpoint_paths = {
            epoch: os.path.join(model_dir, f"model-{epoch:05d}")
            for epoch in range(1, cfg.model.epochs + 1, cfg.validate_every_n_epochs)
        }
        # Add the final checkpoint - catches final epoch != % checkpoint frequency
        # This will be removed if it doesn't exist.
        checkpoint_paths["FINAL"] = os.path.join(model_dir, "model-FINAL")

    # Filter the checkpoints for only checkpoints that exist
    checkpoint_paths = {e: p for e, p in checkpoint_paths.items() if os.path.exists(p + ".index")}

    if cfg.do_validation:
        run_inference('valid',
                      graph_data,
                      checkpoint_paths,
                      cfg,
                      optimizer_options=optimizer_options,
                      losses=losses,
                      loss_weights=loss_weights,
                      metrics=metrics,
                      tmpdir=model_dir)

    if cfg.do_clean_training:
        # force keras eval mode as ogb can crash when forming huge prediction output
        run_inference('clean_train',
                      graph_data,
                      checkpoint_paths,
                      cfg,
                      eval_mode='keras',
                      optimizer_options=optimizer_options,
                      losses=losses,
                      loss_weights=loss_weights,
                      metrics=metrics,
                      tmpdir=model_dir)

    wandb.finish()
    if not cfg.checkpoint_dir:
        logging.info("Removing temporary checkpoint directory.")
        shutil.rmtree(model_dir)


if __name__ == '__main__':
    main()