notebook_utils.py

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

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

import tensorflow as tf
from tensorflow.python.ipu import optimizers
import wandb

import utils
import xpu
from custom_callbacks import CheckpointCallback
from data_utils.input_spec import create_inputs_from_features
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)
from utils import (ThroughputCallback, convert_loss_and_metric_reductions_to_fp32, get_optimizer, str_dtype_to_tf_dtype, set_random_seeds)


def predict(preprocessed_dataset, checkpoint_path, fold, cfg):
    """
    Function to run inference and give predictions of molecule samples.
    Args:
        preprocessed_dataset: dataset object, preprocessed dataset with all the features needed for inference.
        checkpoint_path: string, path to the checkpoint.
        fold: string, name of the fold, choose from ["valid", "test-dev" and "test-challenge"].
        cfg: config object, configurations of the model and inference.
    """
    set_random_seeds()
    tf.keras.mixed_precision.set_global_policy(cfg.model.dtype)

    print("Creating input specification...")
    input_spec = create_inputs_from_features(dataset=preprocessed_dataset, cfg=cfg, fold="test-dev")

    losses, loss_weights = get_loss_functions(preprocessed_dataset, cfg)
    metrics = get_metrics(preprocessed_dataset.denormalize, cfg)
    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
    )

    print("Configuring the IPUs...")
    strategy = xpu.configure_and_get_strategy(num_replicas=1,
                                              num_ipus_per_replica=1,
                                              stochastic_rounding=False,
                                              cfg=cfg)

    with strategy.scope():
        print("Creating TensorFlow dataset from preprocessed dataset...")
        batch_generator, ground_truth_and_masks = get_tf_dataset(preprocessed_dataset=preprocessed_dataset,
                                                                 split_name=fold,
                                                                 shuffle=False,
                                                                 options=cfg,
                                                                 input_spec=input_spec)
        ds = batch_generator.get_tf_dataset()
        ground_truth, include_mask = ground_truth_and_masks
        ground_truth = ground_truth[include_mask]

        print("Constructing the model...")
        model = create_model(batch_generator, preprocessed_dataset, cfg, input_spec=input_spec)
        model.compile(optimizer=get_optimizer(**optimizer_options),
                      loss=losses,
                      loss_weights=loss_weights,
                      weighted_metrics=metrics,
                      steps_per_execution=batch_generator.batches_per_epoch)
        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 checkpoint_path is not None:
            print(f"Loading the checkpoint from {checkpoint_path}...")
            model.load_weights(checkpoint_path).expect_partial()

        print(f"Running `model.predict` to generate predictions...")

        prediction = model.predict(ds, steps=batch_generator.batches_per_epoch)

        if isinstance(prediction, list) and len(prediction) > 1:
            prediction = prediction[0]
        prediction = prediction.squeeze()

        prediction = preprocessed_dataset.denormalize(prediction)

        if len(include_mask) > len(prediction):
            include_mask = include_mask[:len(prediction)]
            ground_truth = ground_truth[:len(prediction)]

        if len(include_mask) > 1:
            include_mask = include_mask.squeeze()

        prediction = prediction[:len(include_mask)][include_mask == 1]

    return prediction, ground_truth


def train(preprocessed_dataset, cfg):
    """
    Utility function to run training on the PCQM4Mv2 dataset from notebook example
    Args:
        preprocessed_dataset: dataset object, preprocessed dataset with all the features needed for inference.
        cfg: config object, configurations of the model and inference.
    """

    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
        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
        return lr
    set_random_seeds()
    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}")

    graph_data = preprocessed_dataset
    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,
    )

    # ------------ 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'))

            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,
                      verbose=1)

    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")}

    return checkpoint_paths