from matplotlib.figure import Figure
from PIL import Image
from pytorch_lightning.loggers import (
LightningLoggerBase,
LoggerCollection,
NeptuneLogger,
TensorBoardLogger,
WandbLogger,
)
from pytorch_lightning.utilities import rank_zero_only
from ride.metrics import FigureDict
from ride.utils.env import RUN_LOGS_PATH
from ride.utils.logging import getLogger, process_rank
logger = getLogger(__name__)
ExperimentLogger = Union[TensorBoardLogger, LoggerCollection, WandbLogger]
ExperimentLoggerCreator = Callable[[str], ExperimentLogger]
def singleton_experiment_logger() -> ExperimentLoggerCreator:
_loggers = {}
def experiment_logger(
name: str = None,
logging_backend: str = "tensorboard",
project_name: str = None,
save_dir=RUN_LOGS_PATH,
) -> ExperimentLogger:
nonlocal _loggers
if logging_backend not in _loggers:
from ptflops import get_model_complexity_info
from supers import supers
from torch import Tensor
from torchmetrics.functional.classification import average_precision
from torchmetrics.functional.classification.confusion_matrix import confusion_matrix
from ride.core import Configs, RideMixin
from ride.utils.logging import getLogger
from ride.utils.utils import merge_dicts, name
ExtendedMetricDict = Dict[str, Union[Tensor, Figure]]
MetricDict = Dict[str, Tensor]
FigureDict = Dict[str, Figure]
StepOutputs = List[Dict[str, Tensor]]
logger = getLogger(__name__, log_once=True)
def sort_out_figures(d: ExtendedMetricDict) -> Tuple[MetricDict, FigureDict]:
mets, figs = {}, {}
for k, v in d.items():
if type(v) == Figure:
figs[k] = v
else:
mets[k] = v
return mets, figs
class OptimisationDirection(Enum):
MIN = "min"
MAX = "max"