base_training.py

"""Main training module for the Joint Learning Super Resolution Face\
 Recognition.
"""
from datetime import datetime
from functools import partial
from pathlib import Path

import tensorflow as tf
from skopt import gp_minimize
from skopt.space import Integer, Real
from skopt.utils import use_named_args
from tensorboard.plugins.hparams import api as hp

from utils.timing import TimingLogger

gpus = tf.config.experimental.list_physical_devices("GPU")
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("set_memory_growth ok!")
    except RuntimeError as e:
        print("set_memory_growth failed!")
        print(str(e))


from tensorflow.keras.mixed_precision import experimental as mixed_precision
from tensorflow_addons.optimizers import AdamW

from models.discriminator import DiscriminatorNetwork
from models.srfr import SRFR
from repositories.casia import CasiaWebface
from services.losses import Loss
from use_cases.train.train_model_joint_learn import TrainModelJointLearnUseCase
from utils.input_data import parseConfigsFile

# Importar Natural DS.
AUTOTUNE = tf.data.experimental.AUTOTUNE


class BaseTraining:
    _CACHE_PATH = Path.cwd().joinpath("temp")

    def __init__(self, logger, strategy, timing: TimingLogger):
        self.logger = logger
        self.strategy = strategy
        self.timing = timing

        if not self._CACHE_PATH.is_dir():
            self._CACHE_PATH.mkdir(parents=True)

    def train(self):
        """Main training function."""
        self.timing.start()

        dimensions = self._create_dimensions()
        hyperparameters = self._create_hyprparameters_domain()
        with tf.summary.create_file_writer(
            str(Path.cwd().joinpath("output", "logs", "hparam_tuning"))
        ).as_default():
            hp.hparams_config(
                hparams=hyperparameters,
                metrics=[hp.Metric("accuracy", display_name="Accuracy")],
            )

        (
            network_settings,
            train_settings,
            preprocess_settings,
        ) = parseConfigsFile(["network", "train", "preprocess"])

        BATCH_SIZE = train_settings["batch_size"] * self.strategy.num_replicas_in_sync

        (
            synthetic_train,
            synthetic_test,
            synthetic_dataset_len,
            synthetic_num_classes,
        ) = self._get_datasets(BATCH_SIZE)

        srfr_model, discriminator_model = self._instantiate_models(
            synthetic_num_classes, network_settings, preprocess_settings
        )

        train_model_use_case = TrainModelJointLearnUseCase(
            self.strategy,
            TimingLogger(),
            self.logger,
            BATCH_SIZE,
            synthetic_dataset_len,
        )

        _training = partial(
            self._fitness_function,
            train_model_use_case=train_model_use_case,
            srfr_model=srfr_model,
            discriminator_model=discriminator_model,
            batch_size=BATCH_SIZE,
            synthetic_train=synthetic_train,
            synthetic_test=synthetic_test,
            num_classes=synthetic_num_classes,
            train_settings=train_settings,
            hparams=hyperparameters,
        )
        _train = use_named_args(dimensions=dimensions)(_training)

        initial_parameters = [
            0.0002,
            2.0e3,
            0.9,
        ]

        search_result = gp_minimize(
            func=_train,
            dimensions=dimensions,
            acq_func="EI",
            n_calls=20,
            x0=initial_parameters,
        )

        self.logger.info(f"Best hyperparameters: {search_result.x}")

    def _fitness_function(
        self,
        learning_rate,
        learning_rate_decay_steps,
        beta_1,
        face_recognition_weight,
        super_resolution_weight,
        perceptual_weight,
        generator_weight,
        l1_weight,
        train_model_use_case: TrainModelJointLearnUseCase,
        srfr_model,
        discriminator_model,
        batch_size,
        synthetic_train,
        synthetic_test,
        num_classes,
        train_settings,
        hparams,
    ):
        (
            HP_LEARNING_RATE,
            HP_LEARNING_RATE_DECAY_STEPS,
            HP_BETA_1,
            HP_FR_WEIGHT,
            HP_SR_WEIGHT,
            HP_PERCEPTUAL_WEIGHT,
            HP_GENERATOR_WEIGHT,
            HP_L1_WEIGHT,
        ) = hparams
        tensorboard_params = {
            HP_LEARNING_RATE: learning_rate,
            HP_LEARNING_RATE_DECAY_STEPS: learning_rate_decay_steps,
            HP_BETA_1: beta_1,
            HP_FR_WEIGHT: face_recognition_weight,
            HP_SR_WEIGHT: super_resolution_weight,
            HP_PERCEPTUAL_WEIGHT: perceptual_weight,
            HP_GENERATOR_WEIGHT: generator_weight,
            HP_L1_WEIGHT: l1_weight,
        }
        learning_rate = self._instantiate_learning_rate(
            learning_rate, learning_rate_decay_steps
        )
        srfr_optimizer, discriminator_optimizer = self._instantiate_optimizers(
            learning_rate, beta_1, train_settings
        )

        summary_writer = self._create_summary_writer()
        metrics = self._instantiate_metrics()

        loss = Loss(
            metrics,
            batch_size,
            summary_writer,
            perceptual_weight,
            generator_weight,
            l1_weight,
            face_recognition_weight,
            super_resolution_weight,
        )

        train_model_use_case.summary_writer = summary_writer

        return -train_model_use_case.execute(
            srfr_model,
            discriminator_model,
            srfr_optimizer,
            discriminator_optimizer,
            synthetic_train,
            synthetic_test,
            num_classes,
            loss,
            tensorboard_params,
        )

    def _create_summary_writer(self):
        with self.strategy.scope():
            current_time = datetime.now().strftime("%Y%m%d-%H%M%S")
            return tf.summary.create_file_writer(
                str(
                    Path.cwd().joinpath("output", "logs", "hparam_tuning", current_time)
                )
            )

    def _instantiate_metrics(self):
        with self.strategy.scope():
            return tf.keras.metrics.CategoricalAccuracy(name="test_crossentropy")

    def _get_datasets(self, batch_size):
        self.logger.info(" -------- Importing Datasets --------")

        casia_dataset = CasiaWebface(self._CACHE_PATH)
        synthetic_train = casia_dataset.get_full_dataset()
        synthetic_train = casia_dataset.augment_dataset(synthetic_train)
        synthetic_train = casia_dataset.normalize_dataset(synthetic_train)

        synthetic_train = synthetic_train.cache(str(self._CACHE_PATH.joinpath("train")))
        synthetic_dataset_len = casia_dataset.get_train_dataset_len()
        synthetic_train = (
            synthetic_train.shuffle(buffer_size=2_048)
            .batch(batch_size, drop_remainder=True)
            .prefetch(1)
        )
        synthetic_train = self.strategy.experimental_distribute_dataset(synthetic_train)

        num_classes = casia_dataset.get_number_of_classes()

        return (
            synthetic_train,
            synthetic_dataset_len,
            num_classes,
        )

    @staticmethod
    def _instantiate_learning_rate(
        learning_rate: float, learning_rate_decay_steps: int
    ):
        return tf.keras.experimental.CosineDecay(
            learning_rate, learning_rate_decay_steps
        )

    def _instantiate_models(
        self,
        strategy,
        synthetic_num_classes,
        network_settings,
        preprocess_settings,
    ):
        self.logger.info(" -------- Creating Models --------")

        with self.strategy.scope():
            srfr_model = SRFR(
                num_filters=network_settings["num_filters"],
                depth=50,
                categories=network_settings["embedding_size"],
                num_gc=network_settings["gc"],
                num_blocks=network_settings["num_blocks"],
                residual_scailing=network_settings["residual_scailing"],
                training=True,
                input_shape=preprocess_settings["image_shape_low_resolution"],
                num_classes_syn=synthetic_num_classes,
            )
            discriminator_model = DiscriminatorNetwork()

        return srfr_model, discriminator_model

    def _instantiate_optimizers(self, learning_rate, beta_1, train_settings):
        self.logger.info(" -------- Creating Optimizers --------")

        with self.strategy.scope():
            srfr_optimizer = AdamW(
                learning_rate=learning_rate,
                beta_1=beta_1,
                beta_2=train_settings["beta_2"],
                weight_decay=train_settings["weight_decay"],
                name="adam_srfr",
            )
            srfr_optimizer = mixed_precision.LossScaleOptimizer(
                srfr_optimizer,
                loss_scale="dynamic",
            )
            discriminator_optimizer = AdamW(
                learning_rate=learning_rate,
                beta_1=beta_1,
                beta_2=train_settings["beta_2"],
                weight_decay=train_settings["weight_decay"],
                name="adam_discriminator",
            )
            discriminator_optimizer = mixed_precision.LossScaleOptimizer(
                discriminator_optimizer, loss_scale="dynamic"
            )

        return (
            srfr_optimizer,
            discriminator_optimizer,
        )

    @staticmethod
    def _create_dimensions():
        return [
            Real(low=1.0e-3, high=7.0e-3, prior="log-uniform", name="learning_rate"),
            Integer(low=1_000, high=6_000, name="learning_rate_decay_steps"),
            Real(low=0.5, high=0.9, prior="log-uniform", name="beta_1"),
            Real(
                low=0.1, high=1.0, prior="log-uniform", name="face_recognition_weight"
            ),
            Real(
                low=0.1, high=1.0, prior="log-uniform", name="super_resolution_weight"
            ),
            Real(
                low=1.0e-3, high=1.0e-2, prior="log-uniform", name="perceptual_weight"
            ),
            Real(low=1.0e-2, high=7.0e-2, prior="log-uniform", name="generator_weight"),
            Real(low=1.0e-2, high=9.0e-2, prior="log-uniform", name="l1_weight"),
        ]

    @staticmethod
    def _create_hyprparameters_domain():
        HP_LEARNING_RATE = hp.HParam(
            "learning_rate",
            hp.RealInterval(1.0e-3, 7.0e-3),
        )
        HP_LEARNING_RATE_DECAY_STEPS = hp.HParam(
            "learning_rate_decay_steps",
            hp.Discrete(range(1_000, 6_000 + 1)),
        )
        HP_BETA_1 = hp.HParam("beta_1", hp.RealInterval(0.5, 0.9))
        HP_FR_WEIGHT = hp.HParam("face_recognition_weight", hp.RealInterval(0.1, 1.0))
        HP_SR_WEIGHT = hp.HParam("super_resolution_weight", hp.RealInterval(0.1, 1.0))
        HP_PERCEPTUAL_WEIGHT = hp.HParam(
            "perceptual_weight", hp.RealInterval(1.0e-3, 1.0e-2)
        )
        HP_GENERATOR_WEIGHT = hp.HParam(
            "generator_weight", hp.RealInterval(1.0e-2, 7.0e-2)
        )
        HP_L1_WEIGHT = hp.HParam("l1_weight", hp.RealInterval(1.0e-2, 9.0e-2))

        return [
            HP_LEARNING_RATE,
            HP_LEARNING_RATE_DECAY_STEPS,
            HP_BETA_1,
            HP_FR_WEIGHT,
            HP_SR_WEIGHT,
            HP_PERCEPTUAL_WEIGHT,
            HP_GENERATOR_WEIGHT,
            HP_L1_WEIGHT,
        ]