utils.py

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

import logging
import time
from typing import Optional

import tensorflow as tf
import wandb
from absl import logging
from tensorflow import keras
import tensorflow.keras.backend as K
from tensorflow.keras import metrics as metrics_mod
import random

import numpy as np
from tensorflow.python import ipu
import xpu


class ThroughputCallback(keras.callbacks.Callback):
    def __init__(self, samples_per_epoch, log_wandb=True):
        super().__init__()
        self.time = 0
        self.samples_per_epoch = samples_per_epoch
        self.log_wandb = log_wandb

    def on_epoch_begin(self, epoch, logs=None):
        self.time = time.time()

    def on_epoch_end(self, epoch, logs=None):
        time_per_epoch = time.time() - self.time
        samples_per_sec = self.samples_per_epoch / time_per_epoch
        if self.log_wandb:
            wandb_logs = {'epoch': epoch, 'throughput': samples_per_sec}
            wandb_logs.update(logs)
            wandb.log(wandb_logs)

        logging.info(f"\nThroughput: {samples_per_sec:.2f} graphs/sec")


def get_optimizer(name="adam",
                  learning_rate=1e-5,
                  l2_regularization=None,
                  dtype="float32",
                  m_dtype=None,
                  v_dtype=None,
                  clip_value=None,
                  loss_scale=1,
                  gradient_accumulation_factor=1,
                  outline_apply_gradients=False,
                  replicas=1):
    def clip_gradients(grads_and_vars):
        return [(tf.clip_by_norm(g, clip_value), v) for g, v in grads_and_vars]

    if clip_value:
        # i think the gradient accumulation factor is incorrect in rescale_gradients.
        # just changing it for the gradient clipping case to ensure i dont break everything else for now
        def rescale_gradients(grads_and_vars):
            return [(g / replicas, v) for g, v in grads_and_vars]

        gradient_transformer = lambda grads_and_vars: clip_gradients(rescale_gradients(grads_and_vars))
    else:

        def rescale_gradients(grads_and_vars):
            return [(g / (gradient_accumulation_factor * replicas), v) for g, v in grads_and_vars]

        gradient_transformer = rescale_gradients

    if name == 'sgd':
        opt_class = add_l2_regularization(tf.keras.optimizers.SGD, l2_regularization)
        opt = opt_class(learning_rate=learning_rate, gradient_transformers=[gradient_transformer])
    elif name == 'tf_adam':
        opt_class = add_l2_regularization(tf.keras.optimizers.Adam, l2_regularization)
        opt = opt_class(learning_rate=learning_rate, gradient_transformers=[gradient_transformer])
    elif name == 'adam':
        opt_class = add_l2_regularization(xpu.AdamIpuOptimizer, l2_regularization)
        opt = opt_class(learning_rate=learning_rate,
                        gradient_transformers=[gradient_transformer],
                        m_dtype=m_dtype,
                        v_dtype=v_dtype,
                        outline_apply_gradients=outline_apply_gradients)
    else:
        raise NotImplementedError(f"Optimizer {name} is not supported.")

    opt = keras.mixed_precision.LossScaleOptimizer(opt, dynamic=False, initial_scale=loss_scale)

    return opt


def add_l2_regularization(optimizer_class, l2_regularization):
    if not l2_regularization:
        return optimizer_class

    class L2Regularizer(optimizer_class):
        def __init__(self, *args, **kwargs):
            super(L2Regularizer, self).__init__(*args, **kwargs)
            self.l2_regularization = l2_regularization

        def _resource_apply_dense(self, grad, var, apply_state):
            return super()._resource_apply_dense(grad + var * self.l2_regularization, var, apply_state)

    return L2Regularizer


def size_hr(num, suffix='B'):
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return "%3.1f %s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)


def print_trainable_variables(model, log_wandb=False):
    logging.info("Trainable Variables:")
    total_parameters = 0
    total_size = 0
    for variable in model.trainable_variables:
        variable_parameters = 1
        for DIM in variable.shape:
            variable_parameters *= DIM
        variable_size = variable_parameters * variable.dtype.size
        logging.info(f"{variable.name}, {variable.shape}, {variable.dtype} ({size_hr(variable_size)})")
        total_parameters += variable_parameters
        total_size += variable_size
    logging.info(f"Total Parameters: {total_parameters:,}  ({size_hr(total_size)})")
    if log_wandb:
        wandb.log({'parameters': total_parameters, 'parameter size (MiB)': total_size / (1024**2)})


def str_dtype_to_tf_dtype(str_dtype):
    return {'float16': tf.float16, 'mixed_float16': tf.float32, 'float32': tf.float32}[str_dtype]


def convert_loss_and_metric_reductions_to_fp32(model):
    model.compiled_loss._loss_metric = metrics_mod.Mean(name='loss', dtype=tf.float32)
    model.compiled_loss._per_output_metrics = [
        metrics_mod.Mean(name=n + '_loss', dtype=tf.float32) for n in model.compiled_loss._output_names
    ]
    metrics = model.compiled_metrics._weighted_metrics
    for k in metrics.keys():
        if not isinstance(metrics[k], list):
            metric_obj = metrics_mod.get(metrics[k])
            if not isinstance(metrics[k], metrics_mod.MeanMetricWrapper):
                metrics[k] = metrics_mod.MeanMetricWrapper(metric_obj, dtype=tf.float32, name=f"{metric_obj.__name__}")
            continue
        new_metrics_list = []
        for m in metrics[k]:
            metric_obj = metrics_mod.get(m)
            if isinstance(metric_obj, metrics_mod.MeanMetricWrapper):
                assert metric_obj.dtype == 'float32', 'metrics must accumulate in float32'
                new_metrics_list += [metric_obj]
            else:
                new_metrics_list += [
                    metrics_mod.MeanMetricWrapper(metric_obj, dtype=tf.float32, name=f"{metric_obj.__name__}")
                ]
        metrics[k] = new_metrics_list
    model.compiled_metrics._weighted_metrics = metrics


def options_validator(cfg):
    """Validate the option combinations that depend on each other. 
    NOTE: This is WIP and not exhaustive. 

    Args:
        cfg : config object from jsonargparse
    """
    # Edges + Noisy Edges
    if cfg.model.use_noisy_edges is True and cfg.model.use_edges is False:
        raise ValueError(f"Cannot use noisy edges without use edges set to True.")

    if cfg.upload_final_ckpt is True and cfg.wandb is False:
        logging.warning(f"`cfg.upload_final_ckpt` is {cfg.upload_final_ckpt} but `cfg.wandb` is {cfg.wandb}." \
                         " Can't upload checkpoint without wandb activated.")


def set_random_seeds(seed=42):
    ipu.utils.reset_ipu_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    tf.random.set_seed(seed)


class Timer:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.t = time.time()

    def __exit__(self, *args, **kwargs):
        elapsed_time = time.time() - self.t
        elapsed_time = time.strftime("%H:%M:%S", time.gmtime(elapsed_time))
        # Timedelta makes for human readable format - not safe for maths operations
        logging.info(f"\U0001F551 Elapsed time step for {self.name}: {elapsed_time} (HH:MM:SS)")