# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import math import os import tempfile import tensorflow as tf from absl import app, flags, logging from ogb.graphproppred import Evaluator import wandb import utils import xpu from data_utils.data_generators import PackedBatchGenerator from model import create_model from data_utils.generated_graph_data import GeneratedGraphData from utils import ThroughputCallback, get_optimizer, str_dtype_to_tf_dtype flags.DEFINE_integer( "micro_batch_size", 8, 'Compute batch size (if using packing this is measured in "packs per batch")' ) flags.DEFINE_integer( "global_batch_size", None, "Global batch size (Includes batches across gradient accumulation factor and replicas)" ) flags.DEFINE_integer("n_nodes_per_pack", 248, 'nodes per "pack"') flags.DEFINE_integer("n_edges_per_pack", 512, 'edges per "pack"') flags.DEFINE_integer("n_graphs_per_pack", 16, 'maximum number of graphs per "pack"') flags.DEFINE_integer("epochs", 100, "maximum number of epochs to run for") flags.DEFINE_float("lr", 2e-5, "learning rate") flags.DEFINE_boolean("cosine_lr", False, "use a cosine lr decay") flags.DEFINE_float("min_lr", 0, "minimum learning rate for the cosine scheduler") flags.DEFINE_float("loss_scaling", 16, "loss scaling factor (to keep gradients representable in IEEE FP16)") flags.DEFINE_enum("adam_m_dtype", "float16", ("float16", "float32"), "dtype for the m part of the adam optimizer") flags.DEFINE_enum("adam_v_dtype", "float32", ("float16", "float32"), "dtype for the v part of the adam optimizer") flags.DEFINE_enum("dataset_name", "ogbg-molhiv", ("ogbg-molhiv",), help="which dataset to use") flags.DEFINE_boolean("generated_data", False, "Use randomly generated data instead of a real dataset.") flags.DEFINE_integer("generated_data_n_nodes", 24, "nodes per graph for the randomly generated dataset") flags.DEFINE_integer("generated_data_n_edges", 50, "edges per graph for the randomly generated dataset") flags.DEFINE_integer( "generated_batches_per_epoch", 2048, "Number of batches per epoch for the randomly generated dataset" ) flags.DEFINE_boolean("do_training", True, "Run training on the dataset") flags.DEFINE_boolean("do_validation", True, "Run validation on the dataset") flags.DEFINE_boolean("do_test", True, "Run test on the dataset") flags.DEFINE_boolean("execution_profile", False, "Create an execution profile in TensorBoard.") flags.DEFINE_boolean("wandb", default=True, help="Enable logging to Weights & Biases") flags.DEFINE_string("checkpoint_path", default=None, help="Path to checkpoint file if skipping training.") flags.DEFINE_integer("checkpoint_every_n_epochs", default=1, help="Create checkpoints every N epochs.") FLAGS = flags.FLAGS class ScaledBinaryCrossentropy(tf.keras.losses.BinaryCrossentropy): def __init__(self, scaling_factor, dtype, **kwargs): super().__init__(**kwargs) self.scaling_factor = tf.cast(scaling_factor, dtype) def __call__(self, y_true, y_pred, sample_weight=None): return self.scaling_factor * super().__call__(y_true, y_pred, sample_weight=sample_weight) def cosine_lr(epoch, lr): completed_fraction = epoch / FLAGS.epochs cosine_decayed = 0.5 * (1.0 + tf.cos(tf.constant(math.pi, dtype=tf.float32) * completed_fraction)) lr = max(cosine_decayed * FLAGS.lr, FLAGS.min_lr) if FLAGS.wandb: wandb.log({"lr": lr}) return lr def main(_): if FLAGS.wandb: wandb.init(entity="sw-apps", project="GIN-test", config=FLAGS.flag_values_dict()) tf.keras.mixed_precision.set_global_policy("float16" if FLAGS.dtype == "float16" else "float32") if FLAGS.generated_data: batch_generator = GeneratedGraphData( FLAGS.micro_batch_size, nodes_per_graph=FLAGS.generated_data_n_nodes, edges_per_graph=FLAGS.generated_data_n_edges, batches_per_epoch=FLAGS.generated_batches_per_epoch, n_graphs_per_pack=FLAGS.n_graphs_per_pack, latent_size=FLAGS.n_latent, ) else: batch_generator = PackedBatchGenerator( dataset_name=FLAGS.dataset_name, n_packs_per_batch=FLAGS.micro_batch_size, fold="train", max_graphs_per_pack=FLAGS.n_graphs_per_pack, max_edges_per_pack=FLAGS.n_edges_per_pack, max_nodes_per_pack=FLAGS.n_nodes_per_pack, n_epochs=FLAGS.epochs, ) ds = batch_generator.get_tf_dataset() if FLAGS.global_batch_size is not None: gradient_accumulation_factor = FLAGS.global_batch_size // (FLAGS.micro_batch_size * FLAGS.replicas) else: gradient_accumulation_factor = 1 steps_per_epoch = batch_generator.batches_per_epoch steps_per_execution_per_replica = steps_per_epoch // FLAGS.replicas steps_per_execution_per_replica = gradient_accumulation_factor * ( steps_per_execution_per_replica // gradient_accumulation_factor ) new_steps_per_epoch = steps_per_execution_per_replica * FLAGS.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=FLAGS.opt.lower(), learning_rate=FLAGS.lr, dtype=str_dtype_to_tf_dtype(FLAGS.dtype), m_dtype=str_dtype_to_tf_dtype(FLAGS.adam_m_dtype), v_dtype=str_dtype_to_tf_dtype(FLAGS.adam_v_dtype), gradient_accumulation_factor=gradient_accumulation_factor, replicas=FLAGS.replicas, ) # Create a temporary directory for the checkpoints with tempfile.TemporaryDirectory() as model_dir: if FLAGS.do_training: strategy_training = xpu.configure_and_get_strategy(FLAGS.replicas) with strategy_training.scope(): model = create_model() utils.print_trainable_variables(model) if FLAGS.loss_scaling > 1.0: losses = ScaledBinaryCrossentropy( scaling_factor=FLAGS.loss_scaling, dtype=str_dtype_to_tf_dtype(FLAGS.dtype) ) else: losses = tf.keras.losses.BinaryCrossentropy() # Weighted metrics are used because of the batch packing weighted_metrics = [tf.keras.metrics.BinaryAccuracy(), tf.keras.metrics.AUC()] 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=FLAGS.wandb, ) ] if FLAGS.cosine_lr: callbacks.append(tf.keras.callbacks.LearningRateScheduler(cosine_lr)) if FLAGS.execution_profile: callbacks.append(tf.keras.callbacks.TensorBoard(profile_batch=[2], log_dir="logs")) if FLAGS.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( tf.keras.callbacks.ModelCheckpoint( model_path, verbose=1, save_best_only=False, save_weights_only=True, period=FLAGS.checkpoint_every_n_epochs, ) ) model.compile( optimizer=get_optimizer(**optimizer_options), loss=losses, weighted_metrics=weighted_metrics, steps_per_execution=steps_per_execution_per_replica, ) # if the total batch size exceeds the compute batch size if xpu.IS_IPU: model.set_gradient_accumulation_options( gradient_accumulation_steps_per_replica=gradient_accumulation_factor, offload_weight_update_variables=False, ) model.fit(ds, steps_per_epoch=steps_per_epoch, epochs=FLAGS.epochs, callbacks=callbacks) if FLAGS.do_validation: logging.info(f"Running validation...") # Use default 1 replica for validation strategy_val = xpu.configure_and_get_strategy(num_replicas=1) with strategy_val.scope(): evaluator = Evaluator(name=FLAGS.dataset_name) if FLAGS.generated_data: valid_gen = batch_generator else: valid_gen = PackedBatchGenerator( dataset_name=FLAGS.dataset_name, n_packs_per_batch=FLAGS.micro_batch_size, fold="valid", max_graphs_per_pack=FLAGS.n_graphs_per_pack, max_edges_per_pack=FLAGS.n_edges_per_pack, max_nodes_per_pack=FLAGS.n_nodes_per_pack, randomize=False, ) val_ds = valid_gen.get_tf_dataset().take(valid_gen.batches_per_epoch).cache().repeat() val_ground_truth, val_include_mask = valid_gen.get_ground_truth_and_masks() val_ground_truth = val_ground_truth[val_include_mask] best_val_auc = 0 model = create_model() model.compile( optimizer=get_optimizer(**optimizer_options), weighted_metrics=weighted_metrics, steps_per_execution=valid_gen.batches_per_epoch, ) if FLAGS.checkpoint_path: checkpoint_paths = {-1: FLAGS.checkpoint_path} logging.info(f"Validating the model with the checkpoint at path {checkpoint_paths}") else: checkpoint_paths = { epoch: os.path.join(model_dir, f"model-{epoch:05d}") for epoch in range(1, FLAGS.epochs + 1) } logging.info(f"Validating over a sweep of checkpoints: {checkpoint_paths}") for epoch, checkpoint_path in checkpoint_paths.items(): model.load_weights(checkpoint_path).expect_partial() prediction = model.predict(val_ds, steps=valid_gen.batches_per_epoch).squeeze() if len(val_include_mask) > len(prediction): val_include_mask = val_include_mask[: len(prediction)] val_ground_truth = val_ground_truth[: len(prediction)] # val_include_mask may be shorter than the predictions — # that is fine (it will just be padding after that point) prediction = prediction[: len(val_include_mask)][val_include_mask.squeeze() == 1] # we will use the official AUC evaluator from the OGB repo, not the keras one result = evaluator.eval({"y_true": val_ground_truth[:, None], "y_pred": prediction[:, None]}) this_auc = result["rocauc"] if FLAGS.wandb: wandb.log({"epoch": epoch, "val_auc": this_auc}) if this_auc > best_val_auc: best_val_auc = this_auc best_model_path = checkpoint_path logging.info(f"{best_val_auc:.3f} is the validation set AUC for the model at {best_model_path}") if FLAGS.do_test: logging.info(f"Running test...") # Use default 1 replica for test strategy_test = xpu.configure_and_get_strategy(num_replicas=1) with strategy_test.scope(): if FLAGS.generated_data: test_gen = batch_generator else: test_gen = PackedBatchGenerator( dataset_name=FLAGS.dataset_name, n_packs_per_batch=FLAGS.micro_batch_size, fold="test", max_graphs_per_pack=FLAGS.n_graphs_per_pack, max_edges_per_pack=FLAGS.n_edges_per_pack, max_nodes_per_pack=FLAGS.n_nodes_per_pack, randomize=False, ) model = create_model() model.compile( optimizer=get_optimizer(**optimizer_options), weighted_metrics=weighted_metrics, steps_per_execution=test_gen.batches_per_epoch, ) test_ds = test_gen.get_tf_dataset().take(test_gen.batches_per_epoch).cache().repeat() test_ground_truth, test_include_mask = test_gen.get_ground_truth_and_masks() test_ground_truth = test_ground_truth[test_include_mask] if FLAGS.checkpoint_path: checkpoint_path = FLAGS.checkpoint_path logging.info(f"Testing the model with the checkpoint at path {checkpoint_path}") else: logging.info(f"Testing the model with the best validation loss (at {best_model_path})") checkpoint_path = best_model_path model.load_weights(checkpoint_path).expect_partial() prediction = model.predict(test_ds, steps=test_gen.batches_per_epoch).squeeze() if len(test_include_mask) > len(prediction): test_include_mask = test_include_mask[: len(prediction)] test_ground_truth = test_ground_truth[: len(prediction)] prediction = prediction[: len(test_include_mask)][test_include_mask.squeeze() == 1] result = evaluator.eval({"y_true": test_ground_truth[:, None], "y_pred": prediction[:, None]}) this_auc = result["rocauc"] logging.info(f"\nTest AUC of {this_auc:.3f}") if FLAGS.wandb: wandb.log({"test_auc": this_auc, "best_val_auc": best_val_auc}) if __name__ == "__main__": app.run(main)