asuiu / keras-generators Goto Github PK

class ModelObject(ABC):
    def __init__(
            self,
            mp: ModelParams,
            model: Model,
            encoders: Dict[str, List[DataEncoder]],
            state_autoclear: int = 128,
    ) -> None:
        """
        :param mp: the model parameters
        :param model: the keras model
        :param encoders: the encoders used to encode/decode the data
        :param state_autoclear: the number of calls to predict/evaluate methods after which the keras session is cleared
        """
        self.mp = mp
        self.model = model
        self.encoders = encoders
        self._state_autoclear = state_autoclear
        self._state_calls = 0

    def _check_clear_state(self):
        self._state_calls += 1  # since it's an atomic operation, it's thread safe due to GIL
        if self._state_calls >= self._state_autoclear:
            K.clear_session()
            logging.debug(f"Cleared keras session after {self._state_calls} calls to predict methods")
            self._state_calls = 0