train.py

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


import argparse
import torch
import poptorch
from pathlib import Path
import time
import numpy as np
import shutil
import warnings
from poptorch import DataLoader
from sklearn.metrics import average_precision_score, roc_auc_score

from tgn_modules import TGN, Data, DataWrapper, init_weights


def build_tgn(
    data: Path,
    dtype: str,
    batch_size: int,
    nodes_size: int,
    edges_size: int,
    dropout: float,
    target: str,
):
    model_dtype = torch.float32 if dtype == "float32" else torch.float16

    train_data = DataWrapper(Data(data, model_dtype, batch_size, nodes_size, edges_size), "train")
    test_data = DataWrapper(Data(data, model_dtype, batch_size, nodes_size, edges_size), "val")

    tgn = TGN(
        num_nodes=9227,
        raw_msg_dim=172,
        memory_dim=100,
        time_dim=100,
        embedding_dim=100,
        dtype=model_dtype,
        dropout=dropout,
        target=target,
    )
    tgn.apply(init_weights)
    return train_data, test_data, tgn


def run_train(model, train_data, optim, target, do_reset) -> float:
    """Trains TGN for one epoch"""
    total_loss = 0
    num_events = 0
    model.train()

    if do_reset:
        model.memory.reset_state()  # Start with a fresh memory.
        if target == "ipu":
            model.copyWeightsToDevice()
    for n, batch in enumerate(train_data):
        if target == "ipu":
            count, loss = model(**batch)
        else:
            optim.zero_grad()
            count, loss = model(**batch)
            loss.backward()
            optim.step()
        total_loss += float(loss) * count
        num_events += count
        model.memory.detach()

    return total_loss / num_events


@torch.no_grad()
def run_test(model, inference_data, target, model_compiled) -> (float, float):
    """Inference over one epoch"""

    model.eval()
    if model_compiled and target == "ipu":
        model.copyWeightsToDevice()
    torch.manual_seed(12345)  # Ensure deterministic sampling across epochs
    aps = 0.0
    aucs = 0.0
    num_events = 0
    for batch in inference_data:
        count, y_true, y_pred = model(**batch)
        aps += count * average_precision_score(y_true, y_pred)
        aucs += count * roc_auc_score(y_true, y_pred)
        num_events += count
    return aps / num_events, aucs / num_events


def _main() -> None:
    parser = argparse.ArgumentParser(
        description="Train the TGN example",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--data", default="data/JODIE", type=Path, help="directory to load/save the data")
    parser.add_argument("-b", "--batch-size", default=40, type=int, help="batch size for training and validation/test")
    parser.add_argument("--nodes-size", default=400, type=int, help="padding for nodes")
    parser.add_argument("--edges-size", default=1200, type=int, help="padding for edges")
    parser.add_argument(
        "-t",
        "--target",
        choices=("cpu", "ipu"),
        default="ipu",
        help="device to run on",
    )
    parser.add_argument(
        "-d",
        "--dtype",
        choices=("float16", "float32"),
        default="float32",
        help="floating point format",
    )
    parser.add_argument(
        "-e",
        "--epochs",
        default=25,
        type=int,
        help="number of epochs to train for",
    )
    parser.add_argument(
        "--validate-every",
        default=1,
        type=int,
        help="run validation every n-th epoch",
    )
    parser.add_argument(
        "--lr",
        default=0.000075,
        type=float,
        help="learning rate",
    )
    parser.add_argument(
        "--dropout",
        default=0.1,
        type=float,
        help="dropout rate in the attention module",
    )
    parser.add_argument(
        "--optimizer",
        choices=("SGD", "Adam"),
        default="Adam",
        help="Optimizer",
    )
    parser.add_argument(
        "-di",
        "--device-iterations",
        default=212,
        type=int,
        help="number of iterations to do on ipu before returning to host",
    )

    args = vars(parser.parse_args())

    epochs = args.pop("epochs")
    validate_every = args.pop("validate_every")
    optim = args.pop("optimizer")
    lr = args.pop("lr")
    device_iterations = args.pop("device_iterations")

    processed_dir = args["data"].joinpath("wikipedia").joinpath("processed")
    # Cleanup artifacts left by previous run
    if processed_dir.exists():
        shutil.rmtree(processed_dir)

    # Build dataloader and model
    train_data, test_data, model = build_tgn(**args)

    num_batches = len(train_data)
    if args["target"] == "cpu":
        device_iterations = int(1)
        warnings.warn(f"device_iterations was set to 1 for training on CPU")
    if not num_batches % device_iterations == 0:
        factors = [x for x in np.arange(1, num_batches + 1) if num_batches % x == 0]
        device_iterations = int(min(factors, key=lambda x: abs(x - device_iterations)))
        warnings.warn(
            f"device_iterations was set to {device_iterations}, the closest " f"factor of the training batch count"
        )

    train_opts = poptorch.Options()
    train_opts.deviceIterations(device_iterations)
    test_opts = poptorch.Options()
    test_opts.deviceIterations(1)

    torch.multiprocessing.set_sharing_strategy("file_system")

    async_options = {
        "sharing_strategy": poptorch.SharingStrategy.SharedMemory,
        "load_indefinitely": True,
        "early_preload": True,
        "buffer_size": 2,
    }
    train_dl = DataLoader(
        options=train_opts,
        dataset=train_data,
        batch_size=1,
        mode=poptorch.DataLoaderMode.Async,
        async_options=async_options,
    )
    test_dl = DataLoader(options=test_opts, dataset=test_data, batch_size=1)

    if args["target"] == "ipu":
        if optim.lower() == "sgd":
            optim = poptorch.optim.SGD(model.parameters(), lr=lr)
        elif optim.lower() == "adam":
            optim = poptorch.optim.AdamW(
                model.parameters(), lr=lr, bias_correction=True, weight_decay=0.0, eps=1e-8, betas=(0.9, 0.999)
            )
        else:
            raise NotImplementedError(f"Optimizer {optim}")
        model_train = poptorch.trainingModel(model, options=train_opts, optimizer=optim)
        model_eval = poptorch.inferenceModel(model, options=test_opts)
    else:
        if optim.lower() == "sgd":
            optim = torch.optim.SGD(model.parameters(), lr=lr)
        elif optim.lower() == "adam":
            optim = torch.optim.Adam(model.parameters(), lr=lr)
        else:
            raise NotImplementedError(f"Optimizer {optim}")
        model_train = model
        model_eval = model

    dataset_size = len(train_dl) * (device_iterations * args["batch_size"])
    test_model_compiled = False

    # Run training
    for epoch in range(1, epochs + 1):
        t0 = time.time()
        loss = run_train(model_train, train_dl, optim, args["target"], do_reset=(epoch > 1))
        duration = time.time() - t0
        tput = dataset_size / duration
        print(f"Epoch {epoch}: Loss {loss:.4f}, Time {duration:.4f}, " f"Throughput: {tput:.4f} samples/s")

        if epoch % validate_every == 0 or epoch == epochs:
            aps, aucs = run_test(model_eval, test_dl, args["target"], model_compiled=test_model_compiled)
            test_model_compiled = True
            print(f"Validation APS {aps:.4f}, AUCS {aucs:.4f}")

    print(f"Training finished \n" f"APS {aps:.4f}, AUCS {aucs:.4f}")

    # Cleanup artifacts created by this run
    shutil.rmtree(processed_dir)


if __name__ == "__main__":
    _main()