easy_tpp/config_factory/model_config.py

from easy_tpp.config_factory.config import Config

from easy_tpp.utils.const import Backend


class TrainerConfig(Config):

    def __init__(self, **kwargs):
        """Initialize the Config class.
        """
        self.seed = kwargs.get('seed', 9899)
        self.gpu = kwargs.get('gpu', -1)
        self.batch_size = kwargs.get('batch_size', 256)
        self.max_epoch = kwargs.get('max_epoch', 10)
        self.shuffle = kwargs.get('shuffle', False)
        self.optimizer = kwargs.get('optimizer', 'adam')
        self.learning_rate = kwargs.get('learning_rate', 1.e-3)
        self.valid_freq = kwargs.get('valid_freq', 1)
        self.use_tfb = kwargs.get('use_tfb', False)
        self.metrics = kwargs.get('metrics', ['acc', 'rmse'])

    def get_yaml_config(self):
        """Return the config in dict (yaml compatible) format.

        Returns:
            dict: config of the trainer specs in dict format.
        """
        return {'seed': self.seed,
                'gpu': self.gpu,
                'batch_size': self.batch_size,
                'max_epoch': self.max_epoch,
                'shuffle': self.shuffle,
                'optimizer': self.optimizer,
                'learning_rate': self.learning_rate,
                'valid_freq': self.valid_freq,
                'use_tfb': self.use_tfb,
                'metrics': self.metrics
                }

    @staticmethod
    def parse_from_yaml_config(yaml_config):
        """Parse from the yaml to generate the config object.

        Args:
            yaml_config (dict): configs from yaml file.

        Returns:
            EasyTPP.TrainerConfig: Config class for trainer specs.
        """
        return TrainerConfig(**yaml_config)

    def copy(self):
        """Copy the config.

        Returns:
            EasyTPP.TrainerConfig: a copy of current config.
        """
        return TrainerConfig(batch_size=self.batch_size,
                             max_epoch=self.max_epoch,
                             shuffle=self.shuffle,
                             optimizer=self.optimizer,
                             learning_rate=self.learning_rate,
                             valid_freq=self.valid_freq,
                             use_tfb=self.use_tfb,
                             metrics=self.metrics
                             )


class ThinningConfig(Config):
    def __init__(self, **kwargs):
        """Initialize the Config class.
        """
        self.num_seq = kwargs.get('num_seq', 10)
        self.num_sample = kwargs.get('num_sample', 1)
        self.num_exp = kwargs.get('num_exp', 500)
        self.look_ahead_time = kwargs.get('look_ahead_time', 10)
        self.patience_counter = kwargs.get('patience_counter', 5)
        self.over_sample_rate = kwargs.get('over_sample_rate', 5)
        self.num_samples_boundary = kwargs.get('num_samples_boundary', 5)
        self.dtime_max = kwargs.get('dtime_max', 5)
        # we pad the sequence at the front only in multi-step generation
        self.num_step_gen = kwargs.get('num_step_gen', 1)

    def get_yaml_config(self):
        """Return the config in dict (yaml compatible) format.

        Returns:
            dict: config of the thinning specs in dict format.
        """
        return {'num_seq': self.num_seq,
                'num_sample': self.num_sample,
                'num_exp': self.num_exp,
                'look_ahead_time': self.look_ahead_time,
                'patience_counter': self.patience_counter,
                'over_sample_rate': self.over_sample_rate,
                'num_samples_boundary': self.num_samples_boundary,
                'dtime_max': self.dtime_max,
                'num_step_gen': self.num_step_gen}

    @staticmethod
    def parse_from_yaml_config(yaml_config):
        """Parse from the yaml to generate the config object.

        Args:
            yaml_config (dict): configs from yaml file.

        Returns:
            EasyTPP.ThinningConfig: Config class for thinning algorithms.
        """
        return ThinningConfig(**yaml_config) if yaml_config is not None else None

    def copy(self):
        """Copy the config.

        Returns:
            EasyTPP.ThinningConfig: a copy of current config.
        """
        return ThinningConfig(num_seq=self.num_seq,
                              num_sample=self.num_sample,
                              num_exp=self.num_exp,
                              look_ahead_time=self.look_ahead_time,
                              patience_counter=self.patience_counter,
                              over_sample_rate=self.over_sample_rate,
                              num_samples_boundary=self.num_samples_boundary,
                              dtime_max=self.dtime_max,
                              num_step_gen=self.num_step_gen)


class BaseConfig(Config):
    def __init__(self, **kwargs):
        """Initialize the Config class.
        """
        self.stage = kwargs.get('stage')
        self.backend = kwargs.get('backend')
        self.dataset_id = kwargs.get('dataset_id')
        self.runner_id = kwargs.get('runner_id')
        self.model_id = kwargs.get('model_id')
        self.exp_id = kwargs.get('exp_id')
        self.base_dir = kwargs.get('base_dir')
        self.specs = kwargs.get('specs', {})
        self.backend = self.set_backend(self.backend)

    @staticmethod
    def set_backend(backend):
        if backend.lower() in ['torch', 'pytorch']:
            return Backend.Torch
        else:
            raise ValueError(
                f"Backend should be 'torch' or 'pytorch', current value: {backend}"
            )

    def get_yaml_config(self):
        """Return the config in dict (yaml compatible) format.

        Returns:
            dict: config of the base config specs in dict format.
        """
        return {'stage': self.stage,
                'backend': str(self.backend),
                'dataset_id': self.dataset_id,
                'runner_id': self.runner_id,
                'model_id': self.model_id,
                'base_dir': self.base_dir,
                'specs': self.specs}

    @staticmethod
    def parse_from_yaml_config(yaml_config):
        """Parse from the yaml to generate the config object.

        Args:
            yaml_config (dict): configs from yaml file.

        Returns:
            BaseConfig: Config class for trainer specs.
        """
        return BaseConfig(**yaml_config)

    def copy(self):
        """Copy the config.

        Returns:
            BaseConfig: a copy of current config.
        """
        return BaseConfig(stage=self.stage,
                          backend=self.backend,
                          dataset_id=self.dataset_id,
                          runner_id=self.runner_id,
                          model_id=self.model_id,
                          base_dir=self.base_dir,
                          specs=self.specs)


class ModelConfig(Config):
    def __init__(self, **kwargs):
        """Initialize the Config class.
        """
        self.rnn_type = kwargs.get('rnn_type', 'LSTM')
        self.hidden_size = kwargs.get('hidden_size', 32)
        self.time_emb_size = kwargs.get('time_emb_size', 16)
        self.num_layers = kwargs.get('num_layers', 2)
        self.num_heads = kwargs.get('num_heads', 2)
        self.sharing_param_layer = kwargs.get('sharing_param_layer', False)
        self.use_mc_samples = kwargs.get('use_mc_samples', True)  # if using MC samples in computing log-likelihood
        self.loss_integral_num_sample_per_step = kwargs.get('loss_integral_num_sample_per_step', 20)  # mc_num_sample_per_step
        self.dropout_rate = kwargs.get('dropout_rate', 0.0)
        self.use_ln = kwargs.get('use_ln', False)
        self.thinning = ThinningConfig.parse_from_yaml_config(kwargs.get('thinning'))
        self.is_training = kwargs.get('training', False)
        self.num_event_types_pad = kwargs.get('num_event_types_pad', None)
        self.num_event_types = kwargs.get('num_event_types', None)
        self.pad_token_id = kwargs.get('event_pad_index', None)
        self.model_id = kwargs.get('model_id', None)
        self.pretrained_model_dir = kwargs.get('pretrained_model_dir', None)
        self.gpu = kwargs.get('gpu', -1)
        self.model_specs = kwargs.get('model_specs', {})

    def get_yaml_config(self):
        """Return the config in dict (yaml compatible) format.

        Returns:
            dict: config of the model config specs in dict format.
        """
        return {'rnn_type': self.rnn_type,
                'hidden_size': self.hidden_size,
                'time_emb_size': self.time_emb_size,
                'num_layers': self.num_layers,
                'sharing_param_layer': self.sharing_param_layer,
                'loss_integral_num_sample_per_step': self.loss_integral_num_sample_per_step,
                'dropout_rate': self.dropout_rate,
                'use_ln': self.use_ln,
                # for some models / cases we may not need to pass thinning config
                # e.g., for intensity-free model
                'thinning': None if self.thinning is None else self.thinning.get_yaml_config(),
                'num_event_types_pad': self.num_event_types_pad,
                'num_event_types': self.num_event_types,
                'event_pad_index': self.pad_token_id,
                'model_id': self.model_id,
                'pretrained_model_dir': self.pretrained_model_dir,
                'gpu': self.gpu,
                'model_specs': self.model_specs}

    @staticmethod
    def parse_from_yaml_config(yaml_config):
        """Parse from the yaml to generate the config object.

        Args:
            yaml_config (dict): configs from yaml file.

        Returns:
            ModelConfig: Config class for trainer specs.
        """
        return ModelConfig(**yaml_config)

    def copy(self):
        """Copy the config.

        Returns:
            ModelConfig: a copy of current config.
        """
        return ModelConfig(rnn_type=self.rnn_type,
                           hidden_size=self.hidden_size,
                           time_emb_size=self.time_emb_size,
                           num_layers=self.num_layers,
                           sharing_param_layer=self.sharing_param_layer,
                           loss_integral_num_sample_per_step=self.loss_integral_num_sample_per_step,
                           dropout_rate=self.dropout_rate,
                           use_ln=self.use_ln,
                           thinning=self.thinning,
                           num_event_types_pad=self.num_event_types_pad,
                           num_event_types=self.num_event_types,
                           event_pad_index=self.pad_token_id,
                           pretrained_model_dir=self.pretrained_model_dir,
                           gpu=self.gpu,
                           model_specs=self.model_specs)