train.py

import atexit
import subprocess
import sys
import numpy as np
import time
import yaml
import torch
from torch.utils.data import DataLoader, SequentialSampler, BatchSampler
import plotting
import pathlib
from test import run_model
import util
from system import DynamicsModelDataset, DynamicsModelFullDataset
from model import load_models, ModelTrainer


def _terminate_server(proc):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()


def create_best_plot(
    model_trainer,
    last_error,
    dataloader,
    model_name,
    model_results_path,
    epoch=0,
    random_state_plot_num=5,
):
    min_val_index = np.argmin(last_error)
    model = model_trainer.models[min_val_index]
    if epoch % 5 == 0:
        random_state_plot_num = 20
    total_error = run_model(
        dataloader,
        model,
        model_trainer.loss_fn,
        random_state_plot_num=random_state_plot_num,
    )
    plotting.plotly_plot_single_model_error(
        model_results_path,
        total_error,
        model_name,
        model,
        prefix="validation_" + str(epoch) + "_" + str(min_val_index),
    )
    result = np.array([total_error["val_loss"]])
    del total_error

    util.clearGPUCache()
    return result


def train_main():
    total_runtime = time.time()
    args = util.loadSharedArgumentsModelLearning("train")

    uid = args["uid"]
    model_config_path = args["model"]
    model_config_file = yaml.load(open(model_config_path, "r"), Loader=yaml.FullLoader)
    model_config_file["dataset"] = args["dataset"]

    models = load_models(args["model"], args)

    if args.get("compile"):
        torch._dynamo.config.recompile_limit = len(models) * 4

    # this is a specific config using the hostname of the computer
    # computer_config = util.getComputerConfig()
    computer_config = {
        "max_batch_size": 10000,
        "train": {"num_workers": 10, "prefetch_factor": 5},
        "test": {"num_workers": 10, "prefetch_factor": 5},
    }

    util.saveGitInformation(model_config_file)

    clear_folder = True
    trained_model_names = []
    server_proc = None

    for model_name, model in models.items():
        # start timing for epoch
        model_start_time = time.time()
        system_type = model.system.name
        system = model.system
        results_dir_path = util.setup_results_dir(
            system_type, uid, "", "train", clear_folder=clear_folder
        )
        clear_folder = False
        trained_model_names.append(model_name)
        plotting.write_server_script(results_dir_path)
        if server_proc is None:
            server_proc = subprocess.Popen(
                [sys.executable, str(results_dir_path / "run.py"), str(args["port"])],
                stdout=subprocess.DEVNULL,
            )
            atexit.register(_terminate_server, server_proc)
            plotting.write_train_landing_page(
                results_dir_path / "index.html",
                list(models.keys()),
                currently_training=model_name,
            )
        # dump the most recent config file based on latest model being run
        yaml.dump(
            model_config_file,
            open(results_dir_path / "model_config.yaml", "w"),
            default_flow_style=False,
            sort_keys=False,
        )

        model_results_path = util.setup_results_dir(
            system_type, uid, model_name, "train", clear_folder=False
        )
        model_results_path.mkdir(parents=True, exist_ok=True)
        yaml.dump(
            model_config_file,
            open(results_dir_path / "model_config.yaml", "w"),
            default_flow_style=False,
            sort_keys=False,
        )

        if "folder_name" in args and args["folder_name"] is not None:
            print(f"creating symbolic link with name ", args["folder_name"])
            sym_path = results_dir_path.parents[0]
            util.createSymLink(
                results_dir_path,
                sym_path / str(args["folder_name"] + "_" + results_dir_path.name),
                relative=True,
            )

        print(f"\n***************************************************")
        print(
            f"model name: {model_name}, "
            # f"size total params: {model.getParamsCount()},"
            # f" recurring params: {model.getRecurringParamsCount()}"
        )
        print(f"using system {model.system.name}")
        print(f"***************************************************")

        traj_length = model.system.traj_length
        epochs = 2
        if "epochs" in args:
            epochs = args["epochs"]

        assert pathlib.Path(args["dataset"]).suffix == ".hdf5"
        dataset_path = pathlib.Path(args["dataset"])

        # print(f"looking for {util.getDatasetName(args['dataset'], system, 'train')} ")

        # if args["dataset"][-5:] == ".hdf5":
        #     dataset_path = args["dataset"]
        # elif util.getDatasetName(args["dataset"], system, "train").exists():
        #     dataset_path = util.getDatasetName(args["dataset"], system, "train")
        # else:
        #     dataset_path = util.getDatasetPath(args["dataset"], system)
        #     system.get_all_bag_data(args)

        assert system.system_type in str(dataset_path)

        # creates dataset classes and returns
        train_dataset = DynamicsModelDataset(
            dataset_path, "train", system.init_length, system.traj_length
        )
        validation_dataset = DynamicsModelFullDataset(
            dataset_path, "validation", system.init_length, system.traj_length
        )
        print(f"val dataset size {len(validation_dataset)}")

        pin_memory = False
        generator = None
        if "cuda" in str(torch.get_default_device()):
            pin_memory = True
            generator = torch.Generator(device="cuda")

        num_workers_val = computer_config["test"]["num_workers"]
        num_workers_train = computer_config["train"]["num_workers"]
        if (
            "training" in model_config_file[model_name].keys()
            and "num_workers" in model_config_file[model_name]["training"].keys()
        ):
            num_workers_val = min(
                num_workers_val,
                model_config_file[model_name]["training"]["num_workers"],
            )
            num_workers_train = min(
                num_workers_train,
                model_config_file[model_name]["training"]["num_workers"],
            )
        prefetch_factor_val = computer_config["test"]["prefetch_factor"]
        prefetch_factor_train = computer_config["train"]["prefetch_factor"]
        if (
            "training" in model_config_file[model_name].keys()
            and "prefetch_factor" in model_config_file[model_name]["training"].keys()
        ):
            prefetch_factor_val = min(
                prefetch_factor_val,
                model_config_file[model_name]["training"]["prefetch_factor"],
            )
            prefetch_factor_train = min(
                prefetch_factor_train,
                model_config_file[model_name]["training"]["prefetch_factor"],
            )

        model_trainer = ModelTrainer(model, model_config_file[model_name])

        train_dataloader = DataLoader(
            train_dataset,
            batch_size=min(model_trainer.batch_size, computer_config["max_batch_size"]),
            num_workers=num_workers_train,
            shuffle=True,
            prefetch_factor=prefetch_factor_train,
            # collate_fn=model_trainer.getTransforms(),
            pin_memory=pin_memory,
            generator=generator,
            persistent_workers=num_workers_train > 0,
            multiprocessing_context="forkserver" if num_workers_train > 0 else None,
        )
        validation_dataloader = DataLoader(
            validation_dataset,
            num_workers=num_workers_val,
            prefetch_factor=prefetch_factor_val,
            batch_sampler=BatchSampler(
                SequentialSampler(validation_dataset),
                # batch_size=min(
                #     model_trainer.batch_size, computer_config["max_batch_size"]
                # ),
                batch_size=10000,
                drop_last=False,
            ),
            # collate_fn=model_trainer.getTransforms(),
            pin_memory=pin_memory,
            multiprocessing_context="forkserver" if num_workers_val > 0 else None,
        )

        print("finished pulling data, going into training")
        print(f"\n***************************************************")
        print(
            f"total train data {len(train_dataset)}, total validation {len(validation_dataset)}"
        )

        # train the model for each bag file
        epoch_stats = {}
        # make a folder for the intermediate models
        model_save_path = model_results_path / "model_store"
        model_save_path.mkdir(parents=True, exist_ok=True)

        training_plotting_path = model_results_path / "model_store"
        training_plotting_path.mkdir(parents=True, exist_ok=True)

        # randomizes all the parallel models
        if not args["finetune"]:
            model_trainer.createRandom()

        # loads a random model to make sure we can load without issues
        model.saveModel(model_save_path, model_name + "_start")
        model.loadModel(model_save_path / (model_name + "_start.pth"))
        model_trainer.saveModel(model_save_path, model_name + "_epoch_start")

        # copies models to the GPU
        if args["gpu"]:
            model_trainer.to("cuda")
        if args["compile"]:
            model_trainer.compile()

        print(f"\n======== START ==========")
        print_iter = 1
        # computes the best model
        last_validation_error = np.array([1])
        # do not need to run this with only a single model, just makes the computation happen twice
        if len(model_trainer.models) > 1 or args["fast_train"]:
            last_validation_error = model_trainer.runValidation(validation_dataloader)
        else:
            last_validation_error = create_best_plot(
                model_trainer,
                last_validation_error,
                validation_dataloader,
                model_name,
                model_results_path,
                epoch=-1,
            )

        print(f"with random model got validation loss:\n{last_validation_error}")
        last_min_validation_error = np.min(last_validation_error)
        if len(model_trainer.models) > 1 and not args["fast_train"]:
            create_best_plot(
                model_trainer,
                last_validation_error,
                validation_dataloader,
                model_name,
                model_results_path,
                epoch=-1,
            )

        # set up dict for loging purposes
        epoch_stats["loss"] = np.zeros((epochs, len(model_trainer.models)))
        epoch_stats["validation_error"] = np.zeros((epochs, len(model_trainer.models)))
        epoch_stats["grad_norm"] = np.zeros((epochs, len(model_trainer.models)))

        model_trainer.reloadOptimizer()

        for epoch in range(epochs):
            if epoch % print_iter == 0:
                print(f"\n======== {epoch+1}/{epochs} ==========")

            # train
            loss = (
                model_trainer.trainModel(train_dataloader, model_save_path)
                .cpu()
                .detach()
                .numpy()
            )
            epoch_stats["loss"][epoch] = loss
            epoch_stats["grad_norm"][epoch] = model_trainer.epoch_grad_norms

            # to validation
            val_error = None
            if len(model_trainer.models) > 1 or args["fast_train"]:
                val_error = model_trainer.runValidation(validation_dataloader)
            else:
                val_error = create_best_plot(
                    model_trainer,
                    last_validation_error,
                    validation_dataloader,
                    model_name,
                    model_results_path,
                    epoch=epoch + 1,
                )

            epoch_stats["validation_error"][epoch] = val_error

            # save intermediate models
            model_trainer.saveModel(
                model_save_path, model_name + "_epoch_" + str(epoch)
            )

            # === Plots the val/loss curve
            plotting.plotly_plot_val_loss_curve(
                epoch_stats, model_results_path / "loss_curve", epoch + 1
            )
            model_trainer.plotParameters(training_plotting_path, epoch + 1)
            plotting.write_model_train_page(model_results_path, model_name)
            # TODO plot the intermediate values to a pdf

            # print user information for debugging
            min_val_index = np.argmin(val_error)
            min_validation_error = val_error[min_val_index]
            if epoch % print_iter == 0:
                np.set_printoptions(precision=4)
                print(f"\non epoch {epoch+1} got loss:\n{loss}")
                print(f"on epoch {epoch+1} got validation loss:\n{val_error}")
                print(
                    f" validation delta to best:     {min_validation_error - last_min_validation_error:+.4f}"
                )
                print(
                    f" validation delta to previous:\n{val_error - last_validation_error}"
                )

            last_min_validation_error = min(
                min_validation_error, last_min_validation_error
            )
            last_validation_error = val_error

            if len(model_trainer.models) > 1 and not args["fast_train"]:
                create_best_plot(
                    model_trainer,
                    last_validation_error,
                    validation_dataloader,
                    model_name,
                    model_results_path,
                    epoch=epoch + 1,
                )

            del val_error, loss

            # TODO fancy stopping of the model learning based of flat validation error

        del train_dataloader
        print(f"\n=========== finished running training =============")

        # === Plots the val/loss curve
        plotting.plotly_plot_val_loss_curve(
            epoch_stats, model_results_path / "loss_curve", epochs
        )
        model_trainer.plotParameters(training_plotting_path, epochs)

        # pick the best model that we have trained
        ind = np.unravel_index(
            np.argmin(epoch_stats["validation_error"], axis=None),
            epoch_stats["validation_error"].shape,
        )
        model_config_file[model_name]["path"] = str(
            model_save_path / (model_name + ".pth")
        )
        # create symoblic link to the model with the lowest validation error
        util.createSymLink(
            model_save_path
            / (
                model_name
                + "_epoch_"
                + str(ind[0])
                + "_model_num_"
                + str(ind[1])
                + ".pth"
            ),
            model_save_path / (model_name + ".pth"),
            relative=True,
        )
        # dump the most recent config file based on latest model being run
        yaml.dump(
            model_config_file,
            open(results_dir_path / "model_config.yaml", "w"),
            default_flow_style=False,
            sort_keys=False,
        )

        # loads the best performing model
        model.loadModel(model_config_file[model_name]["path"])

        if args["gpu"]:
            model.to("cuda")
        if args["compile"]:
            model.compile()

        # runs the validation dataset to get error over entire trajectory
        total_error = run_model(
            validation_dataloader,
            model,
            model_trainer.loss_fn,
            random_state_plot_num=20,
            num_bins=5000,
            T_delta=0.25,
        )
        plotting.plotly_plot_single_model_error(
            model_results_path, total_error, model_name, model, prefix="validation"
        )
        del validation_dataset
        del validation_dataloader
        del total_error

        # plots errors for the training dataset
        train_dataset = DynamicsModelFullDataset(
            dataset_path, "train", system.init_length, system.traj_length
        )
        train_dataloader = DataLoader(
            train_dataset,
            batch_sampler=BatchSampler(
                SequentialSampler(train_dataset),
                batch_size=10000,
                drop_last=False,
            ),
            num_workers=num_workers_val,
            prefetch_factor=prefetch_factor_val,
            persistent_workers=True,
            # collate_fn=model_trainer.getTransforms(),
            pin_memory=pin_memory,
            multiprocessing_context="forkserver" if num_workers_val > 0 else None,
        )
        total_error = run_model(
            train_dataloader,
            model,
            model_trainer.loss_fn,
            num_bins=5000,
            T_delta=0.25,
            random_state_plot_num=20,
        )
        plotting.plotly_plot_single_model_error(
            model_results_path, total_error, model_name, model, prefix="train"
        )

        plotting.write_model_train_page(model_results_path, model_name)

        _model_keys = list(models.keys())
        _next_model = (
            _model_keys[_model_keys.index(model_name) + 1]
            if _model_keys.index(model_name) + 1 < len(_model_keys)
            else None
        )
        plotting.write_train_landing_page(
            results_dir_path / "index.html",
            _model_keys,
            currently_training=_next_model,
        )

        # del train_dataset
        # del train_dataloader
        models[model_name] = None
        # del total_error

        # model_trainer.GPUClean()
        del model_trainer
        del model

        print(f"***************************************************")
        print(f"finished running {model_name} {time.time() - model_start_time:.2f}s")
        print(f"***************************************************")

    del models
    print(f"took {time.time() - total_runtime:.2f}s to complete entire run")
    if not args["no_pause"]:
        print(
            f"\nResults: http://localhost:{args['port']}  (serving {results_dir_path})"
        )
        print("Press Ctrl+C to stop the server.")
        try:
            server_proc.wait()
        except KeyboardInterrupt:
            pass  # atexit handler will terminate the server on exit

    # TODO should now run the test.py on the final results we have


if __name__ == "__main__":
    train_main()