power_cogs.utils.logging — power_cogs 0.1.0 documentation

# Source code for power_cogs.utils.logging

    import os
    from functools import wraps

    import tensorflow as tf

    def safe(f):
        """wrap function with try / catch
        Args:
            f ([type]): function to wrap
        """

        @wraps(f)
        def new_func(*args, **kwargs):
            try:
                f(*args, **kwargs)
            except Exception as e:
                print("{} failed with {}".format(f.__name__, e))

        return new_func

    class TensorboardLogger:
        def __init__(self, tensorboard_log_path: str = "./tensorboard_logs"):
            self.tensorboard_run_path = os.path.abspath(
                os.path.join(tensorboard_log_path, "default")
            )
            self.tensorboard_writer = self.setup_tensorboard(self.tensorboard_run_path)

        def setup_tensorboard(self, tensorboard_run_path: str):
            tensorboard_writer = tf.summary.create_file_writer(tensorboard_run_path)
            tensorboard_writer.set_as_default()
            return tensorboard_writer

        @safe
        def log_scalar(self, data, name="", step=0) -> None:
            if self.tensorboard_writer is not None:
                self.tensorboard_writer.set_as_default()
                tf.summary.scalar(name="1/" + name, data=data, step=step)

        @safe
        def log_text(self, data, name="", step=0, description="") -> None:
            if self.tensorboard_writer is not None:
                tf.summary.text(name=name, data=data, step=0, description=description)

    def doublewrap(f):
        """
        from https://stackoverflow.com/questions/653368/how-to-create-a-python-decorator-that-can-be-used-either-with-or-without-paramet
        answer by @bj0
        a decorator decorator, allowing the decorator to be used as:
        @decorator(with, arguments, and=kwargs)
        or
        @decorator
        """

        @wraps(f)
        def new_dec(*args, **kwargs):
            if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
                # actual decorated function
                return f(args[0])
            else:
                # decorator arguments
                return lambda realf: f(realf, *args, **kwargs)

        return new_dec

    @doublewrap
    def tensorboard(cls):
        class TensorboardWrapper(cls):
            def __init__(self, *args, **kwargs):
                self.logger = TensorboardLogger()
                self.__name__ = cls.__name__
                super().__init__(*args, **kwargs)

            def __repr__(self):
                return repr(cls)

            def log_scalar(self, *args, **kwargs):
                self.logger.log_scalar(*args, **kwargs)

            def log_text(self, *args, **kwargs):
                self.logger.log_text(*args, **kwargs)

        return TensorboardWrapper