164 lines
5.2 KiB
Python
164 lines
5.2 KiB
Python
from argparse import Namespace
|
|
from typing import Optional, Dict, Any, Union
|
|
|
|
try:
|
|
from test_tube import Experiment
|
|
except ImportError: # pragma: no-cover
|
|
raise ImportError('You want to use `test_tube` logger which is not installed yet,' # pragma: no-cover
|
|
' install it with `pip install test-tube`.')
|
|
|
|
from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_only
|
|
|
|
|
|
class TestTubeLogger(LightningLoggerBase):
|
|
r"""
|
|
Log to local file system in TensorBoard format but using a nicer folder structure.
|
|
(see `full docs <https://williamfalcon.github.io/test-tube>`_).
|
|
"""
|
|
|
|
__test__ = False
|
|
|
|
def __init__(
|
|
self, save_dir: str, name: str = "default", description: Optional[str] = None,
|
|
debug: bool = False, version: Optional[int] = None, create_git_tag: bool = False
|
|
):
|
|
r"""
|
|
|
|
Example
|
|
----------
|
|
|
|
.. code-block:: python
|
|
|
|
logger = TestTubeLogger("tt_logs", name="my_exp_name")
|
|
trainer = Trainer(logger=logger)
|
|
trainer.train(model)
|
|
|
|
Use the logger anywhere in you LightningModule as follows:
|
|
|
|
.. code-block:: python
|
|
|
|
def train_step(...):
|
|
# example
|
|
self.logger.experiment.whatever_method_summary_writer_supports(...)
|
|
|
|
def any_lightning_module_function_or_hook(...):
|
|
self.logger.experiment.add_histogram(...)
|
|
|
|
Args:
|
|
save_dir (str): Save directory
|
|
name (str): Experiment name. Defaults to "default".
|
|
description (str): A short snippet about this experiment
|
|
debug (bool): If True, it doesn't log anything
|
|
version (int): Experiment version. If version is not specified the logger inspects the save
|
|
directory for existing versions, then automatically assigns the next available version.
|
|
create_git_tag (bool): If True creates a git tag to save the code used in this experiment
|
|
|
|
"""
|
|
super().__init__()
|
|
self.save_dir = save_dir
|
|
self._name = name
|
|
self.description = description
|
|
self.debug = debug
|
|
self._version = version
|
|
self.create_git_tag = create_git_tag
|
|
self._experiment = None
|
|
|
|
@property
|
|
def experiment(self) -> Experiment:
|
|
r"""
|
|
|
|
Actual test-tube object. To use test-tube features do the following.
|
|
|
|
Example::
|
|
|
|
self.logger.experiment.some_test_tube_function()
|
|
|
|
"""
|
|
|
|
if self._experiment is not None:
|
|
return self._experiment
|
|
|
|
self._experiment = Experiment(
|
|
save_dir=self.save_dir,
|
|
name=self._name,
|
|
debug=self.debug,
|
|
version=self.version,
|
|
description=self.description,
|
|
create_git_tag=self.create_git_tag,
|
|
rank=self.rank,
|
|
)
|
|
return self._experiment
|
|
|
|
@rank_zero_only
|
|
def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
|
|
# TODO: HACK figure out where this is being set to true
|
|
self.experiment.debug = self.debug
|
|
params = self._convert_params(params)
|
|
params = self._flatten_dict(params)
|
|
self.experiment.argparse(Namespace(**params))
|
|
|
|
@rank_zero_only
|
|
def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
|
|
# TODO: HACK figure out where this is being set to true
|
|
self.experiment.debug = self.debug
|
|
self.experiment.log(metrics, global_step=step)
|
|
|
|
@rank_zero_only
|
|
def save(self) -> None:
|
|
# TODO: HACK figure out where this is being set to true
|
|
self.experiment.debug = self.debug
|
|
self.experiment.save()
|
|
|
|
@rank_zero_only
|
|
def finalize(self, status: str) -> None:
|
|
# TODO: HACK figure out where this is being set to true
|
|
self.experiment.debug = self.debug
|
|
self.save()
|
|
self.close()
|
|
|
|
@rank_zero_only
|
|
def close(self) -> None:
|
|
# TODO: HACK figure out where this is being set to true
|
|
self.experiment.debug = self.debug
|
|
if not self.debug:
|
|
exp = self.experiment
|
|
exp.close()
|
|
|
|
@property
|
|
def rank(self) -> int:
|
|
return self._rank
|
|
|
|
@rank.setter
|
|
def rank(self, value: int) -> None:
|
|
self._rank = value
|
|
if self._experiment is not None:
|
|
self.experiment.rank = value
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
if self._experiment is None:
|
|
return self._name
|
|
else:
|
|
return self.experiment.name
|
|
|
|
@property
|
|
def version(self) -> int:
|
|
if self._experiment is None:
|
|
return self._version
|
|
else:
|
|
return self.experiment.version
|
|
|
|
# Test tube experiments are not pickleable, so we need to override a few
|
|
# methods to get DDP working. See
|
|
# https://docs.python.org/3/library/pickle.html#handling-stateful-objects
|
|
# for more info.
|
|
def __getstate__(self) -> Dict[Any, Any]:
|
|
state = self.__dict__.copy()
|
|
state["_experiment"] = self.experiment.get_meta_copy()
|
|
return state
|
|
|
|
def __setstate__(self, state: Dict[Any, Any]):
|
|
self._experiment = state["_experiment"].get_non_ddp_exp()
|
|
del state["_experiment"]
|
|
self.__dict__.update(state)
|