train.py

import json
import math
import random
import shutil
import traceback
from enum import Enum
from functools import wraps
from typing import *

import os
import sys
import mltk
import tensorkit as tk
import numpy as np
import torch
import click
from tensorkit import tensor as T
from tensorkit.examples import utils
from tensorkit.train import Checkpoint

from tracegnn.data import *
from tracegnn.models.trace_vae.evaluation import *
from tracegnn.models.trace_vae.graph_utils import *
from tracegnn.models.trace_vae.tensor_utils import *
from tracegnn.models.trace_vae.types import *
from tracegnn.models.trace_vae.model import *
from tracegnn.models.trace_vae.dataset import *
from tracegnn.utils import *


class NANLossError(Exception):

    def __init__(self, epoch):
        super().__init__(epoch)

    @property
    def epoch(self) -> Optional[int]:
        return self.args[0]

    def __str__(self):
        return f'NaN loss encountered at epoch {self.epoch}'


class OptimizerType(str, Enum):
    ADAM = 'adam'
    RMSPROP = 'rmsprop'


class ExpConfig(mltk.Config):
    model: TraceVAEConfig = TraceVAEConfig()
    device: Optional[str] = 'cpu'
    seed: Optional[int] = 0

    class train(mltk.Config):
        max_epoch: int = 60
        struct_pretrain_epochs: Optional[int] = 40  # number of epochs to pre-train the struct_vae
        ckpt_epoch_freq: Optional[int] = 5
        test_epoch_freq: Optional[int] = 5
        latency_hist_epoch_freq: Optional[int] = 10
        latency_std_hist_epoch_freq: Optional[int] = 5

        use_early_stopping: bool = False
        val_epoch_freq: Optional[int] = 2

        kl_beta: float = 1.0
        warm_up_epochs: Optional[int] = None  # number of epochs to warm-up the prior (KLD)

        l2_reg: float = 0.0001
        z_unit_ball_reg: Optional[float] = None
        z2_unit_ball_reg: Optional[float] = None

        init_batch_size: int = 64
        batch_size: int = 64
        val_batch_size: int = 64

        optimizer: OptimizerType = OptimizerType.RMSPROP
        initial_lr: float = 0.001
        lr_anneal_ratio: float = 0.1
        lr_anneal_epochs: int = 30
        clip_norm: Optional[float] = None
        global_clip_norm: Optional[float] = 10  # important for numerical stability

        test_n_z: int = 10
        num_plot_samples: int = 20

    class test(mltk.Config):
        batch_size: int = 64
        eval_n_z: int = 10
        use_biased: bool = True
        latency_log_prob_weight: bool = True
        clip_nll: Optional[float] = 100_000

    class report(mltk.Config):
        html_ext: str = '.html.gz'

    class dataset(mltk.Config):
        root_dir: str = os.path.abspath('./data/processed')


def main(exp: mltk.Experiment[ExpConfig]):
    # config
    config = exp.config

    # set random seed to encourage reproducibility (does it really work?)
    if config.seed is not None:
        T.random.set_deterministic(True)
        T.random.seed(config.seed)
        np.random.seed(config.seed)
        random.seed(config.seed)

    # Load data
    id_manager = TraceGraphIDManager(os.path.join(config.dataset.root_dir, 'id_manager'))
    latency_range = TraceGraphLatencyRangeFile(os.path.join(config.dataset.root_dir, 'id_manager'))

    train_db = TraceGraphDB(BytesSqliteDB(os.path.join(config.dataset.root_dir, 'processed', 'train')))
    val_db = TraceGraphDB(BytesSqliteDB(os.path.join(config.dataset.root_dir, 'processed', 'val')))
    test_db = TraceGraphDB(
        BytesMultiDB(
            BytesSqliteDB(os.path.join(config.dataset.root_dir, 'processed', 'test')),
            BytesSqliteDB(os.path.join(config.dataset.root_dir, 'processed', 'test-drop')),
            BytesSqliteDB(os.path.join(config.dataset.root_dir, 'processed', 'test-latency')),
        )
    )
    train_stream = TraceGraphDataStream(
        train_db, id_manager=id_manager, batch_size=config.train.batch_size,
        shuffle=True, skip_incomplete=False,
    )
    val_stream = TraceGraphDataStream(
        val_db, id_manager=id_manager, batch_size=config.train.val_batch_size,
        shuffle=False, skip_incomplete=False,
    )
    test_stream = TraceGraphDataStream(
        test_db, id_manager=id_manager, batch_size=config.test.batch_size,
        shuffle=False, skip_incomplete=False,
    )

    utils.print_experiment_summary(
        exp,
        train_data=train_stream,
        val_data=val_stream,
        test_data=test_stream
    )
    print('Train Data:', train_db)
    print('Val Data:', val_db)
    print('Test Data:', test_db)

    # build the network
    vae: TraceVAE = TraceVAE(
        config.model,
        id_manager.num_operations,
    )
    vae = vae.to(T.current_device())
    params, param_names = utils.get_params_and_names(vae)
    utils.print_parameters_summary(params, param_names)
    print('')
    mltk.print_with_time('Network constructed.')

    # define the training method for a certain model part
    def train_part(params, start_epoch, max_epoch, latency_only, do_final_eval):
        # util to ensure all installed hooks will only run within this context
        in_context = [True]

        def F(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                if in_context[0]:
                    return func(*args, **kwargs)
            return wrapper

        # the train procedure
        try:
            # buffer to collect stds of each p(latency|z)
            latency_std = {}
            for key in ('train', 'val', 'test_normal', 'test_drop', 'test_latency'):
                latency_std[key] = ArrayBuffer(81920)

            def should_collect_latency_std():
                return (
                    config.train.latency_std_hist_epoch_freq and
                    loop.epoch % config.train.latency_std_hist_epoch_freq == 0
                )

            def clear_std_buf():
                for buf in latency_std.values():
                    buf.clear()

            # the initialization function
            def initialize():
                G = TraceGraphBatch(
                    id_manager=id_manager,
                    latency_range=latency_range,
                    trace_graphs=train_db.sample_n(config.train.init_batch_size),
                )
                chain = vae.q(G).chain(
                    vae.p,
                    G=G,
                )
                loss = chain.vi.training.sgvb(reduction='mean')
                mltk.print_with_time(f'Network initialized: loss = {T.to_numpy(loss)}')

            # the train functions
            def on_train_epoch_begin():
                # set train mode
                if latency_only:
                    tk.layers.set_eval_mode(vae)
                    tk.layers.set_train_mode(vae.latency_vae)
                else:
                    tk.layers.set_train_mode(vae)

                # clear std buffer
                clear_std_buf()

            def train_step(trace_graphs):
                G = TraceGraphBatch(
                    id_manager=id_manager,
                    latency_range=latency_range,
                    trace_graphs=trace_graphs,
                )
                chain = vae.q(G).chain(
                    vae.p,
                    G=G,
                )

                # collect the latency std
                if should_collect_latency_std():
                    collect_latency_std(latency_std['train'], chain)

                # collect the log likelihoods
                p_obs = []
                p_latent = []
                q_latent = []
                for name in chain.p:
                    if name in chain.q:
                        q_latent.append(chain.q[name].log_prob())
                        p_latent.append(chain.p[name].log_prob())
                    else:
                        # print(name, chain.p[name].log_prob().mean())
                        p_obs.append(chain.p[name].log_prob())

                # get E[log p(x|z)] and KLD[q(z|x)||p(z)]
                recons = T.reduce_mean(T.add_n(p_obs))
                kl = T.reduce_mean(T.add_n(q_latent) - T.add_n(p_latent))

                # KL beta
                beta = config.train.kl_beta
                if config.train.warm_up_epochs and loop.epoch < config.train.warm_up_epochs:
                    beta = beta * (loop.epoch / config.train.warm_up_epochs)
                loss = beta * kl - recons

                # l2 regularization
                if config.train.l2_reg:
                    l2_params = []
                    for p, n in zip(params, param_names):
                        if 'bias' not in n:
                            l2_params.append(p)
                    loss = loss + config.train.l2_reg * T.nn.l2_regularization(l2_params)

                # unit ball regularization
                def add_unit_ball_reg(l, t, reg):
                    if reg is not None:
                        ball_mean, ball_var = get_moments(t, axis=[-1])
                        l = l + reg * (
                            T.reduce_mean(ball_mean ** 2) +
                            T.reduce_mean((ball_var - 1) ** 2)
                        )
                    return l

                loss = add_unit_ball_reg(loss, chain.q['z'].tensor, config.train.z_unit_ball_reg)
                if 'z2' in chain.q:
                    loss = add_unit_ball_reg(loss, chain.q['z2'].tensor, config.train.z2_unit_ball_reg)

                # check and return the metrics
                loss_val = T.to_numpy(loss)
                if math.isnan(loss_val):
                    raise NANLossError(loop.epoch)

                return {'loss': loss, 'recons': recons, 'kl': kl}

            # the validation function
            def validate():
                tk.layers.set_eval_mode(vae)

                def val_step(trace_graphs):
                    with T.no_grad():
                        G = TraceGraphBatch(
                            id_manager=id_manager,
                            latency_range=latency_range,
                            trace_graphs=trace_graphs,
                        )
                        chain = vae.q(G).chain(
                            vae.p,
                            G=G,
                        )
                        # collect the latency std
                        if should_collect_latency_std():
                            collect_latency_std(latency_std['val'], chain)
                        loss = chain.vi.training.sgvb()
                        return {'loss': T.to_numpy(T.reduce_mean(loss))}

                val_loop = loop.validation()
                result_dict = val_loop.run(val_step, val_stream)
                result_dict = {
                    f'val_{k}': v
                    for k, v in result_dict.items()
                }
                summary_cb.update_metrics(result_dict)

            # the evaluation function
            def evaluate(n_z, eval_loop, eval_stream, epoch, use_embeddings=False,
                         plot_latency_hist=False):
                # latency_hist_file
                latency_hist_file = None
                if plot_latency_hist:
                    latency_hist_file = exp.make_parent(f'./plotting/latency-sample/{epoch}.jpg')

                # do evaluation
                tk.layers.set_eval_mode(vae)
                with T.no_grad():
                    kw = {}
                    if should_collect_latency_std():
                        kw['latency_std_dict_out'] = latency_std
                        kw['latency_dict_prefix'] = 'test_'
                    result_dict = do_evaluate_nll(
                        test_stream=eval_stream,
                        vae=vae,
                        id_manager=id_manager,
                        latency_range=latency_range,
                        n_z=n_z,
                        use_biased=config.test.use_biased,
                        latency_log_prob_weight=config.test.latency_log_prob_weight,
                        test_loop=eval_loop,
                        summary_writer=summary_cb,
                        clip_nll=config.test.clip_nll,
                        use_embeddings=use_embeddings,
                        latency_hist_file=latency_hist_file,
                        **kw,
                    )

                with open(exp.make_parent(f'./result/test-anomaly/{epoch}.json'), 'w', encoding='utf-8') as f:
                    f.write(json.dumps(result_dict))
                eval_loop.add_metrics(**result_dict)

            def save_model(epoch=None):
                epoch = epoch or loop.epoch
                torch.save(vae.state_dict(), exp.make_parent(f'models/{epoch}.pt'))

            # final evaluation
            if do_final_eval:
                tk.layers.set_eval_mode(vae)

                # save the final model
                save_model('final')

                clear_std_buf()
                evaluate(
                    n_z=config.test.eval_n_z,
                    eval_loop=mltk.TestLoop(),
                    eval_stream=test_stream,
                    epoch='final',
                    use_embeddings=True,
                    plot_latency_hist=True,
                )

            else:
                # set train mode at the beginning of each epoch
                loop.on_epoch_begin.do(F(on_train_epoch_begin))

                # the optimizer and learning rate scheduler
                if config.train.optimizer == OptimizerType.ADAM:
                    optimizer = tk.optim.Adam(params)
                elif config.train.optimizer == OptimizerType.RMSPROP:
                    optimizer = tk.optim.RMSprop(params)

                def update_lr():
                    n_cycles = int(
                        loop.epoch //  # (loop.epoch - start_epoch) //
                        config.train.lr_anneal_epochs
                    )
                    lr_discount = config.train.lr_anneal_ratio ** n_cycles
                    optimizer.set_lr(config.train.initial_lr * lr_discount)

                update_lr()
                loop.on_epoch_end.do(F(update_lr))

                # install the validation function and early-stopping
                if config.train.val_epoch_freq:
                    loop.run_after_every(
                        F(validate),
                        epochs=config.train.val_epoch_freq,
                    )

                # install the evaluation function during training
                if config.train.test_epoch_freq:
                    loop.run_after_every(
                        F(lambda: evaluate(
                            n_z=config.train.test_n_z,
                            eval_loop=loop.test(),
                            eval_stream=test_stream,
                            epoch=loop.epoch,
                            plot_latency_hist=(
                                config.train.latency_hist_epoch_freq and
                                loop.epoch % config.train.latency_hist_epoch_freq == 0
                            )
                        )),
                        epochs=config.train.test_epoch_freq,
                    )

                # install the plot and sample functions during training
                def after_epoch():
                    save_model()
                loop.run_after_every(F(after_epoch), epochs=1)

                # train the model
                tk.layers.set_eval_mode(vae)
                on_train_epoch_begin()
                initialize()
                utils.fit_model(
                    loop=loop,
                    optimizer=optimizer,
                    fn=train_step,
                    stream=train_stream,
                    clip_norm=config.train.clip_norm,
                    global_clip_norm=config.train.global_clip_norm,
                    # pass to `loop.run()`
                    limit=max_epoch,
                )
        finally:
            in_context = [False]

    # the train loop
    loop = mltk.TrainLoop(max_epoch=config.train.max_epoch)

    # checkpoint
    ckpt = Checkpoint(vae=vae)
    loop.add_callback(mltk.callbacks.AutoCheckpoint(
        ckpt,
        root_dir=exp.make_dirs('./checkpoint'),
        epoch_freq=config.train.ckpt_epoch_freq,
        max_checkpoints_to_keep=10,
    ))

    # early-stopping
    if config.train.val_epoch_freq and config.train.use_early_stopping:
        loop.add_callback(mltk.callbacks.EarlyStopping(
            checkpoint=ckpt,
            root_dir=exp.abspath('./early-stopping'),
            metric_name='val_loss',
        ))

    # the summary writer
    summary_cb = SummaryCallback(summary_dir=exp.abspath('./summary'))
    loop.add_callback(summary_cb)

    # pre-train the struct_vae
    try:
        with loop:
            start_epoch = 1
            part_params = params
            latency_only = False

            if (config.model.arch == TraceVAEArch.DEFAULT) and config.train.struct_pretrain_epochs:
                # train struct_vae first
                print(f'Start to train vae with {len(part_params)} params ...')
                train_part(
                    list(part_params),
                    start_epoch=start_epoch,
                    max_epoch=config.train.struct_pretrain_epochs,
                    latency_only=latency_only,
                    do_final_eval=False,
                )

                # train latency_vae next
                part_params = [
                    p for n, p in zip(param_names, params)
                    if n.startswith('latency_vae')
                ]
                start_epoch = config.train.struct_pretrain_epochs + 1
                latency_only = True
                print(f'Start to train latency_vae with {len(part_params)} params ...')

            train_part(
                part_params,
                start_epoch=start_epoch,
                max_epoch=config.train.max_epoch,
                latency_only=latency_only,
                do_final_eval=False,
            )

        # do final evaluation
        train_part(
            [],
            start_epoch=-1,
            max_epoch=-1,
            latency_only=False,
            do_final_eval=True,
        )

    except KeyboardInterrupt:
        print(
            'Train interrupted, press Ctrl+C again to skip the final test ...',
            file=sys.stderr,
        )


if __name__ == '__main__':
    with mltk.Experiment(ExpConfig) as exp:
        config = exp.config
        device = config.device or T.first_gpu_device()
        with T.use_device(device):
            retrial = 0
            while True:
                try:
                    main(exp)
                except NANLossError as ex:
                    if ex.epoch != 1 or retrial >= 10:
                        raise
                    retrial += 1
                    print(
                        f'\n'
                        f'Restart the experiment for the {retrial}-th time '
                        f'due to NaN loss at epoch {ex.epoch}.\n',
                        file=sys.stderr
                    )
                    if ex.epoch == 1:
                        for name in ['checkpoint', 'early-stopping', 'models',
                                     'plotting', 'summary']:
                            path = exp.abspath(name)
                            if os.path.isdir(name):
                                shutil.rmtree(path)
                else:
                    break