<filename>src/backend/training/training.py
import json
import logging
import math
import os

import os.path
import pickle
import subprocess
from typing import List

from keras.callbacks import History

from networks import train_nn
from game.models import ValuePolicyModel
from generators import Experience
from helpers import get_timestamp
from utils.plots import generate_all_plots

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
logging.basicConfig(filename="../runtime.log", level=logging.DEBUG, filemode="w",
                    format="%(asctime)s - %(process)d - %(levelname)s - %(message)s",
                    datefmt="%H:%M:%S")


class ModelTrainer:

    def __init__(self, training_kwargs: dict = None, game_kwargs: dict = None,
                 mcts_kwargs: dict = None, nn_kwargs: dict = None):
        self.training_kwargs = training_kwargs or {}
        self.game_kwargs = game_kwargs or {}
        self.mcts_kwargs = mcts_kwargs or {}
        self.nn_kwargs = nn_kwargs or {}

        self.training_iterations = self.training_kwargs.get("TRAINING_ITERATIONS", 1)
        self.starting_iteration = self.training_kwargs.get("TRAINING_START_ITERATION", 0)
        self.input_data_base_dir = self.training_kwargs.get("INPUT_DIR",
                                                            "data/training/")  # TODO unify save paths
        self.output_base_dir = self.training_kwargs.get("OUTPUT_DIR", "data")
        self.max_processes = self.training_kwargs.get("MAX_PROCESSES", 1)

        # Additional data added to training data from previous iterations
        self.training_previous = self.training_kwargs.get("TRAINING_PREVIOUS", 0)

        self.games_per_iteration = self.training_kwargs.get("GAMES_PER_ITERATION", 50)
        self.rollouts_per_game = self.training_kwargs.get("ROLLOUTS_PER_GAME", 1000)
        self.terminate_counter = self.training_kwargs.get("TERMINATE_COUNTER", 50)

        model_base_name = self.training_kwargs.get("MODEL_BASE_NAME", "model")
        self.model = ValuePolicyModel()
        self.model.set_name(model_base_name)

    def __call__(self, generate_data=True, generate_plots=False):
        history, best_checkpoint_filepath = None, None

        for iteration_id in range(
                self.starting_iteration, self.starting_iteration + self.training_iterations):

            if generate_data:
                self.generate_data(iteration=iteration_id)

            training_data = []
            for data_iteration_id in range(
                    max(0, iteration_id - self.training_previous), iteration_id+1):
                input_dir = self.get_input_dir(iteration=iteration_id)
                training_data.extend(self.load_training_data(input_dir))

            self.model.load(filename=iteration_id)
            history, checkpoint_filepath = train_nn(
                training_data, self.model, iteration=iteration_id, **self.nn_kwargs)

            self.model.load_checkpoint(checkpoint_filepath)
            self.model.save(filename=str(iteration_id + 1))
            self.save_history(history=history, iteration=iteration_id)

            if generate_plots:
                source = f"../data/history/{self.model.name}"
                destination = f"../data/plots/{self.model.name}"
                generate_all_plots(source, destination)

        return history, checkpoint_filepath

    @staticmethod
    def load_training_data(input_dir: str) -> List["Experience"]:
        """
        Collects all records saved inside input dir and merges them into a single input data.

        :param input_dir: Directory from which data will be loaded.

        :return: List of merged data records.
        """
        training_data = []
        for filename in os.listdir(input_dir):
            filepath = os.path.join(input_dir, filename)

            with open(filepath, "rb") as file_:
                training_data.extend(pickle.load(file_))

        return training_data

    def generate_data(self, iteration: int = -1):
        games = math.ceil(self.games_per_iteration / self.max_processes)

        processes = []
        logging.debug("Starting Data generation")
        for idx in range(self.max_processes):
            logging.debug(os.listdir())
            processes.append(
                subprocess.Popen(
                    [
                        "python", "training/generators.py",
                        str(games),
                        str(iteration),
                        str(self.terminate_counter),
                        str(self.rollouts_per_game),
                        str(self.model.name)
                    ])
            )

        for process in processes:
            logging.debug(f"Starting data generation using process {process.pid}")
            process.wait()
            logging.debug(f"Completed data generation using process {process.pid}")

    def get_input_dir(self, iteration: int) -> str:
        """
        Returns data input directory given iteration number.

        :param iteration: Iteration ID defining data folder.

        :return: Relative path to directory from base dir.
        """
        dir_name = f"iteration_{iteration}"
        return os.path.join(self.input_data_base_dir, self.model.name, dir_name)

    def save_history(self, history: History, iteration: int) -> None:
        """
        Saves keras History to disc.

        :param history: History to save.
        :param iteration: Number of iteration at which save occurs.
        :return: None
        """
        base_path = os.path.join(self.output_base_dir, "history", self.model.name)
        os.makedirs(base_path, exist_ok=True)

        filename = f"{get_timestamp()}_{iteration}.json"
        filepath = os.path.join(self.output_base_dir, "history", self.model.name,
                                filename)  # TODO unify save paths
        with open(filepath, "w") as file_:
            logging.info(f"Saving training history to file {filename}.")
            json.dump(history.history, file_)


if __name__ == '__main__':
    training_kwargs = {
        "TRAINING_ITERATIONS": 15,
        "TRAINING_START_ITERATION": 20,
        "TRAINING_PREVIOUS": 10,
        "INPUT_DIR": "data/training/",
        "OUTPUT_DIR": "data/",
        "MAX_PROCESSES": 8,
        "MODEL_BASE_NAME": "first_model",
        "GAMES_PER_ITERATION": 200,
        "ROLLOUTS_PER_GAME": 300,
    }
    game_kwargs = {}
    mcts_kwargs = {}
    nn_kwargs = {
        "EPOCHS": 15,
        "VALIDATION_SPLIT": 0.10,
    }

    model_trainer = ModelTrainer(training_kwargs, game_kwargs, mcts_kwargs, nn_kwargs)
    model_trainer(generate_data=True)

<filename>src/backend/training/training.py
import json
import logging
import math
import os

import os.path
import pickle
import subprocess
from typing import List

from keras.callbacks import History

from networks import train_nn
from game.models import ValuePolicyModel
from generators import Experience
from helpers import get_timestamp
from utils.plots import generate_all_plots

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
logging.basicConfig(filename="../runtime.log", level=logging.DEBUG, filemode="w",
                    format="%(asctime)s - %(process)d - %(levelname)s - %(message)s",
                    datefmt="%H:%M:%S")


class ModelTrainer:

    def __init__(self, training_kwargs: dict = None, game_kwargs: dict = None,
                 mcts_kwargs: dict = None, nn_kwargs: dict = None):
        self.training_kwargs = training_kwargs or {}
        self.game_kwargs = game_kwargs or {}
        self.mcts_kwargs = mcts_kwargs or {}
        self.nn_kwargs = nn_kwargs or {}

        self.training_iterations = self.training_kwargs.get("TRAINING_ITERATIONS", 1)
        self.starting_iteration = self.training_kwargs.get("TRAINING_START_ITERATION", 0)
        self.input_data_base_dir = self.training_kwargs.get("INPUT_DIR",
                                                            "data/training/")  # TODO unify save paths
        self.output_base_dir = self.training_kwargs.get("OUTPUT_DIR", "data")
        self.max_processes = self.training_kwargs.get("MAX_PROCESSES", 1)

        # Additional data added to training data from previous iterations
        self.training_previous = self.training_kwargs.get("TRAINING_PREVIOUS", 0)

        self.games_per_iteration = self.training_kwargs.get("GAMES_PER_ITERATION", 50)
        self.rollouts_per_game = self.training_kwargs.get("ROLLOUTS_PER_GAME", 1000)
        self.terminate_counter = self.training_kwargs.get("TERMINATE_COUNTER", 50)

        model_base_name = self.training_kwargs.get("MODEL_BASE_NAME", "model")
        self.model = ValuePolicyModel()
        self.model.set_name(model_base_name)

    def __call__(self, generate_data=True, generate_plots=False):
        history, best_checkpoint_filepath = None, None

        for iteration_id in range(
                self.starting_iteration, self.starting_iteration + self.training_iterations):

            if generate_data:
                self.generate_data(iteration=iteration_id)

            training_data = []
            for data_iteration_id in range(
                    max(0, iteration_id - self.training_previous), iteration_id+1):
                input_dir = self.get_input_dir(iteration=iteration_id)
                training_data.extend(self.load_training_data(input_dir))

            self.model.load(filename=iteration_id)
            history, checkpoint_filepath = train_nn(
                training_data, self.model, iteration=iteration_id, **self.nn_kwargs)

            self.model.load_checkpoint(checkpoint_filepath)
            self.model.save(filename=str(iteration_id + 1))
            self.save_history(history=history, iteration=iteration_id)

            if generate_plots:
                source = f"../data/history/{self.model.name}"
                destination = f"../data/plots/{self.model.name}"
                generate_all_plots(source, destination)

        return history, checkpoint_filepath

    @staticmethod
    def load_training_data(input_dir: str) -> List["Experience"]:
        """
        Collects all records saved inside input dir and merges them into a single input data.

        :param input_dir: Directory from which data will be loaded.

        :return: List of merged data records.
        """
        training_data = []
        for filename in os.listdir(input_dir):
            filepath = os.path.join(input_dir, filename)

            with open(filepath, "rb") as file_:
                training_data.extend(pickle.load(file_))

        return training_data

    def generate_data(self, iteration: int = -1):
        games = math.ceil(self.games_per_iteration / self.max_processes)

        processes = []
        logging.debug("Starting Data generation")
        for idx in range(self.max_processes):
            logging.debug(os.listdir())
            processes.append(
                subprocess.Popen(
                    [
                        "python", "training/generators.py",
                        str(games),
                        str(iteration),
                        str(self.terminate_counter),
                        str(self.rollouts_per_game),
                        str(self.model.name)
                    ])
            )

        for process in processes:
            logging.debug(f"Starting data generation using process {process.pid}")
            process.wait()
            logging.debug(f"Completed data generation using process {process.pid}")

    def get_input_dir(self, iteration: int) -> str:
        """
        Returns data input directory given iteration number.

        :param iteration: Iteration ID defining data folder.

        :return: Relative path to directory from base dir.
        """
        dir_name = f"iteration_{iteration}"
        return os.path.join(self.input_data_base_dir, self.model.name, dir_name)

    def save_history(self, history: History, iteration: int) -> None:
        """
        Saves keras History to disc.

        :param history: History to save.
        :param iteration: Number of iteration at which save occurs.
        :return: None
        """
        base_path = os.path.join(self.output_base_dir, "history", self.model.name)
        os.makedirs(base_path, exist_ok=True)

        filename = f"{get_timestamp()}_{iteration}.json"
        filepath = os.path.join(self.output_base_dir, "history", self.model.name,
                                filename)  # TODO unify save paths
        with open(filepath, "w") as file_:
            logging.info(f"Saving training history to file {filename}.")
            json.dump(history.history, file_)


if __name__ == '__main__':
    training_kwargs = {
        "TRAINING_ITERATIONS": 15,
        "TRAINING_START_ITERATION": 20,
        "TRAINING_PREVIOUS": 10,
        "INPUT_DIR": "data/training/",
        "OUTPUT_DIR": "data/",
        "MAX_PROCESSES": 8,
        "MODEL_BASE_NAME": "first_model",
        "GAMES_PER_ITERATION": 200,
        "ROLLOUTS_PER_GAME": 300,
    }
    game_kwargs = {}
    mcts_kwargs = {}
    nn_kwargs = {
        "EPOCHS": 15,
        "VALIDATION_SPLIT": 0.10,
    }

    model_trainer = ModelTrainer(training_kwargs, game_kwargs, mcts_kwargs, nn_kwargs)
    model_trainer(generate_data=True)

# TODO unify save paths # Additional data added to training data from previous iterations Collects all records saved inside input dir and merges them into a single input data.          :param input_dir: Directory from which data will be loaded.          :return: List of merged data records. Returns data input directory given iteration number.          :param iteration: Iteration ID defining data folder.          :return: Relative path to directory from base dir. Saves keras History to disc.          :param history: History to save.         :param iteration: Number of iteration at which save occurs.         :return: None # TODO unify save paths