399 lines
14 KiB
Python
399 lines
14 KiB
Python
"""
|
|
TRAINS
|
|
------
|
|
"""
|
|
from argparse import Namespace
|
|
from os import environ
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL.Image import Image
|
|
|
|
try:
|
|
import trains
|
|
from trains import Task
|
|
_TRAINS_AVAILABLE = True
|
|
except ImportError: # pragma: no-cover
|
|
trains = None
|
|
Task = None
|
|
_TRAINS_AVAILABLE = False
|
|
raise ImportError('You want to use `TRAINS` logger which is not installed yet,' # pragma: no-cover
|
|
' install it with `pip install trains`.')
|
|
|
|
from pytorch_lightning import _logger as log
|
|
from pytorch_lightning.loggers.base import LightningLoggerBase
|
|
from pytorch_lightning.utilities import rank_zero_only
|
|
|
|
|
|
class TrainsLogger(LightningLoggerBase):
|
|
"""
|
|
Log using `allegro.ai TRAINS <https://github.com/allegroai/trains>`_. Install it with pip:
|
|
|
|
.. code-block:: bash
|
|
|
|
pip install trains
|
|
|
|
Example:
|
|
>>> from pytorch_lightning import Trainer
|
|
>>> from pytorch_lightning.loggers import TrainsLogger
|
|
>>> trains_logger = TrainsLogger(
|
|
... project_name='pytorch lightning',
|
|
... task_name='default',
|
|
... output_uri='.',
|
|
... ) # doctest: +ELLIPSIS
|
|
TRAINS Task: ...
|
|
TRAINS results page: ...
|
|
>>> trainer = Trainer(logger=trains_logger)
|
|
|
|
Use the logger anywhere in your :class:`~pytorch_lightning.core.lightning.LightningModule` as follows:
|
|
|
|
>>> from pytorch_lightning import LightningModule
|
|
>>> class LitModel(LightningModule):
|
|
... def training_step(self, batch, batch_idx):
|
|
... # example
|
|
... self.logger.experiment.whatever_trains_supports(...)
|
|
...
|
|
... def any_lightning_module_function_or_hook(self):
|
|
... self.logger.experiment.whatever_trains_supports(...)
|
|
|
|
Args:
|
|
project_name: The name of the experiment's project. Defaults to ``None``.
|
|
task_name: The name of the experiment. Defaults to ``None``.
|
|
task_type: The name of the experiment. Defaults to ``'training'``.
|
|
reuse_last_task_id: Start with the previously used task id. Defaults to ``True``.
|
|
output_uri: Default location for output models. Defaults to ``None``.
|
|
auto_connect_arg_parser: Automatically grab the :class:`~argparse.ArgumentParser`
|
|
and connect it with the task. Defaults to ``True``.
|
|
auto_connect_frameworks: If ``True``, automatically patch to trains backend. Defaults to ``True``.
|
|
auto_resource_monitoring: If ``True``, machine vitals will be
|
|
sent along side the task scalars. Defaults to ``True``.
|
|
|
|
Examples:
|
|
>>> logger = TrainsLogger("pytorch lightning", "default", output_uri=".") # doctest: +ELLIPSIS
|
|
TRAINS Task: ...
|
|
TRAINS results page: ...
|
|
>>> logger.log_metrics({"val_loss": 1.23}, step=0)
|
|
>>> logger.log_text("sample test")
|
|
sample test
|
|
>>> import numpy as np
|
|
>>> logger.log_artifact("confusion matrix", np.ones((2, 3)))
|
|
>>> logger.log_image("passed", "Image 1", np.random.randint(0, 255, (200, 150, 3), dtype=np.uint8))
|
|
"""
|
|
|
|
_bypass = None
|
|
|
|
def __init__(
|
|
self,
|
|
project_name: Optional[str] = None,
|
|
task_name: Optional[str] = None,
|
|
task_type: str = 'training',
|
|
reuse_last_task_id: bool = True,
|
|
output_uri: Optional[str] = None,
|
|
auto_connect_arg_parser: bool = True,
|
|
auto_connect_frameworks: bool = True,
|
|
auto_resource_monitoring: bool = True
|
|
) -> None:
|
|
if not _TRAINS_AVAILABLE:
|
|
raise ImportError('You want to use `test_tube` logger which is not installed yet,'
|
|
' install it with `pip install test-tube`.')
|
|
super().__init__()
|
|
if self.bypass_mode():
|
|
self._trains = None
|
|
print('TRAINS Task: running in bypass mode')
|
|
print('TRAINS results page: disabled')
|
|
|
|
class _TaskStub(object):
|
|
def __call__(self, *args, **kwargs):
|
|
return self
|
|
|
|
def __getattr__(self, attr):
|
|
if attr in ('name', 'id'):
|
|
return ''
|
|
return self
|
|
|
|
def __setattr__(self, attr, val):
|
|
pass
|
|
|
|
self._trains = _TaskStub()
|
|
else:
|
|
self._trains = Task.init(
|
|
project_name=project_name,
|
|
task_name=task_name,
|
|
task_type=task_type,
|
|
reuse_last_task_id=reuse_last_task_id,
|
|
output_uri=output_uri,
|
|
auto_connect_arg_parser=auto_connect_arg_parser,
|
|
auto_connect_frameworks=auto_connect_frameworks,
|
|
auto_resource_monitoring=auto_resource_monitoring
|
|
)
|
|
|
|
@property
|
|
def experiment(self) -> Task:
|
|
r"""
|
|
Actual TRAINS object. To use TRAINS features in your
|
|
:class:`~pytorch_lightning.core.lightning.LightningModule` do the following.
|
|
|
|
Example::
|
|
|
|
self.logger.experiment.some_trains_function()
|
|
|
|
"""
|
|
return self._trains
|
|
|
|
@property
|
|
def id(self) -> Union[str, None]:
|
|
"""
|
|
ID is a uuid (string) representing this specific experiment in the entire system.
|
|
"""
|
|
if not self._trains:
|
|
return None
|
|
|
|
return self._trains.id
|
|
|
|
@rank_zero_only
|
|
def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
|
|
"""
|
|
Log hyperparameters (numeric values) in TRAINS experiments.
|
|
|
|
Args:
|
|
params: The hyperparameters that passed through the model.
|
|
"""
|
|
if not self._trains:
|
|
return
|
|
if not params:
|
|
return
|
|
|
|
params = self._convert_params(params)
|
|
params = self._flatten_dict(params)
|
|
self._trains.connect(params)
|
|
|
|
@rank_zero_only
|
|
def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
|
|
"""
|
|
Log metrics (numeric values) in TRAINS experiments.
|
|
This method will be called by Trainer.
|
|
|
|
Args:
|
|
metrics: The dictionary of the metrics.
|
|
If the key contains "/", it will be split by the delimiter,
|
|
then the elements will be logged as "title" and "series" respectively.
|
|
step: Step number at which the metrics should be recorded. Defaults to ``None``.
|
|
"""
|
|
if not self._trains:
|
|
return
|
|
|
|
if not step:
|
|
step = self._trains.get_last_iteration()
|
|
|
|
for k, v in metrics.items():
|
|
if isinstance(v, str):
|
|
log.warning("Discarding metric with string value {}={}".format(k, v))
|
|
continue
|
|
if isinstance(v, torch.Tensor):
|
|
v = v.item()
|
|
parts = k.split('/')
|
|
if len(parts) <= 1:
|
|
series = title = k
|
|
else:
|
|
title = parts[0]
|
|
series = '/'.join(parts[1:])
|
|
self._trains.get_logger().report_scalar(
|
|
title=title, series=series, value=v, iteration=step)
|
|
|
|
@rank_zero_only
|
|
def log_metric(self, title: str, series: str, value: float, step: Optional[int] = None) -> None:
|
|
"""
|
|
Log metrics (numeric values) in TRAINS experiments.
|
|
This method will be called by the users.
|
|
|
|
Args:
|
|
title: The title of the graph to log, e.g. loss, accuracy.
|
|
series: The series name in the graph, e.g. classification, localization.
|
|
value: The value to log.
|
|
step: Step number at which the metrics should be recorded. Defaults to ``None``.
|
|
"""
|
|
if not self._trains:
|
|
return
|
|
|
|
if not step:
|
|
step = self._trains.get_last_iteration()
|
|
|
|
if isinstance(value, torch.Tensor):
|
|
value = value.item()
|
|
self._trains.get_logger().report_scalar(
|
|
title=title, series=series, value=value, iteration=step)
|
|
|
|
@rank_zero_only
|
|
def log_text(self, text: str) -> None:
|
|
"""Log console text data in TRAINS experiment.
|
|
|
|
Args:
|
|
text: The value of the log (data-point).
|
|
"""
|
|
if self.bypass_mode():
|
|
print(text)
|
|
return
|
|
|
|
if not self._trains:
|
|
return
|
|
|
|
self._trains.get_logger().report_text(text)
|
|
|
|
@rank_zero_only
|
|
def log_image(
|
|
self, title: str, series: str,
|
|
image: Union[str, np.ndarray, Image, torch.Tensor],
|
|
step: Optional[int] = None) -> None:
|
|
"""
|
|
Log Debug image in TRAINS experiment
|
|
|
|
Args:
|
|
title: The title of the debug image, i.e. "failed", "passed".
|
|
series: The series name of the debug image, i.e. "Image 0", "Image 1".
|
|
image: Debug image to log. If :class:`numpy.ndarray` or :class:`torch.Tensor`,
|
|
the image is assumed to be the following:
|
|
|
|
- shape: CHW
|
|
- color space: RGB
|
|
- value range: [0., 1.] (float) or [0, 255] (uint8)
|
|
|
|
step: Step number at which the metrics should be recorded. Defaults to None.
|
|
"""
|
|
if not self._trains:
|
|
return
|
|
|
|
if not step:
|
|
step = self._trains.get_last_iteration()
|
|
|
|
if isinstance(image, str):
|
|
self._trains.get_logger().report_image(
|
|
title=title, series=series, local_path=image, iteration=step)
|
|
else:
|
|
if isinstance(image, torch.Tensor):
|
|
image = image.cpu().numpy()
|
|
if isinstance(image, np.ndarray):
|
|
image = image.transpose(1, 2, 0)
|
|
self._trains.get_logger().report_image(
|
|
title=title, series=series, image=image, iteration=step)
|
|
|
|
@rank_zero_only
|
|
def log_artifact(
|
|
self, name: str,
|
|
artifact: Union[str, Path, Dict[str, Any], np.ndarray, Image],
|
|
metadata: Optional[Dict[str, Any]] = None, delete_after_upload: bool = False) -> None:
|
|
"""
|
|
Save an artifact (file/object) in TRAINS experiment storage.
|
|
|
|
Args:
|
|
name: Artifact name. Notice! it will override the previous artifact
|
|
if the name already exists.
|
|
artifact: Artifact object to upload. Currently supports:
|
|
|
|
- string / :class:`pathlib.Path` are treated as path to artifact file to upload
|
|
If a wildcard or a folder is passed, a zip file containing the
|
|
local files will be created and uploaded.
|
|
- dict will be stored as .json file and uploaded
|
|
- :class:`pandas.DataFrame` will be stored as .csv.gz (compressed CSV file) and uploaded
|
|
- :class:`numpy.ndarray` will be stored as .npz and uploaded
|
|
- :class:`PIL.Image.Image` will be stored to .png file and uploaded
|
|
|
|
metadata:
|
|
Simple key/value dictionary to store on the artifact. Defaults to ``None``.
|
|
delete_after_upload:
|
|
If ``True``, the local artifact will be deleted (only applies if ``artifact`` is a
|
|
local file). Defaults to ``False``.
|
|
"""
|
|
if not self._trains:
|
|
return
|
|
|
|
self._trains.upload_artifact(
|
|
name=name, artifact_object=artifact, metadata=metadata,
|
|
delete_after_upload=delete_after_upload
|
|
)
|
|
|
|
@rank_zero_only
|
|
def finalize(self, status: str = None) -> None:
|
|
# super().finalize(status)
|
|
if self.bypass_mode() or not self._trains:
|
|
return
|
|
|
|
self._trains.close()
|
|
self._trains = None
|
|
|
|
@property
|
|
def name(self) -> Union[str, None]:
|
|
"""
|
|
Name is a human readable non-unique name (str) of the experiment.
|
|
"""
|
|
if not self._trains:
|
|
return ''
|
|
|
|
return self._trains.name
|
|
|
|
@property
|
|
def version(self) -> Union[str, None]:
|
|
if not self._trains:
|
|
return None
|
|
|
|
return self._trains.id
|
|
|
|
@classmethod
|
|
def set_credentials(cls, api_host: str = None, web_host: str = None, files_host: str = None,
|
|
key: str = None, secret: str = None) -> None:
|
|
"""
|
|
Set new default TRAINS-server host and credentials.
|
|
These configurations could be overridden by either OS environment variables
|
|
or trains.conf configuration file.
|
|
|
|
Note:
|
|
Credentials need to be set *prior* to Logger initialization.
|
|
|
|
Args:
|
|
api_host: Trains API server url, example: ``host='http://localhost:8008'``
|
|
web_host: Trains WEB server url, example: ``host='http://localhost:8080'``
|
|
files_host: Trains Files server url, example: ``host='http://localhost:8081'``
|
|
key: user key/secret pair, example: ``key='thisisakey123'``
|
|
secret: user key/secret pair, example: ``secret='thisisseceret123'``
|
|
"""
|
|
Task.set_credentials(api_host=api_host, web_host=web_host, files_host=files_host,
|
|
key=key, secret=secret)
|
|
|
|
@classmethod
|
|
def set_bypass_mode(cls, bypass: bool) -> None:
|
|
"""
|
|
Will bypass all outside communication, and will drop all logs.
|
|
Should only be used in "standalone mode", when there is no access to the *trains-server*.
|
|
|
|
Args:
|
|
bypass: If ``True``, all outside communication is skipped.
|
|
"""
|
|
cls._bypass = bypass
|
|
|
|
@classmethod
|
|
def bypass_mode(cls) -> bool:
|
|
"""
|
|
Returns the bypass mode state.
|
|
|
|
Note:
|
|
`GITHUB_ACTIONS` env will automatically set bypass_mode to ``True``
|
|
unless overridden specifically with ``TrainsLogger.set_bypass_mode(False)``.
|
|
|
|
Return:
|
|
If True, all outside communication is skipped.
|
|
"""
|
|
return cls._bypass if cls._bypass is not None else bool(environ.get('CI'))
|
|
|
|
def __getstate__(self) -> Union[str, None]:
|
|
if self.bypass_mode() or not self._trains:
|
|
return ''
|
|
|
|
return self._trains.id
|
|
|
|
def __setstate__(self, state: str) -> None:
|
|
self._rank = 0
|
|
self._trains = None
|
|
if state:
|
|
self._trains = Task.get_task(task_id=state)
|