mirror of
https://github.com/wassname/pytorch-lightning.git
synced 2026-10-29 13:30:57 +08:00
Nic Eggert
480eed5cb6
* Implement generic loggers for experiment tracking * Add tests for loggers * Get model tests passing * Test and fix logger pickling * Expand pickle test and fix bug * Missed exp -> logger conversion * Remove commented code * Add docstrings * Update logging docs * Add mlflow to test requirements * Make linter happy * Fix mlflow timestamp * Update Logging.md * Update test_models.py * Update test_models.py * Update test_models.py * Update properties.md * Fix tests * Line length
45 lines
1.3 KiB
Python
45 lines
1.3 KiB
Python
from time import time
|
|
from logging import getLogger
|
|
|
|
import mlflow
|
|
|
|
from .base import LightningLoggerBase, rank_zero_only
|
|
|
|
logger = getLogger(__name__)
|
|
|
|
|
|
class MLFlowLogger(LightningLoggerBase):
|
|
|
|
def __init__(self, experiment_name, tracking_uri=None):
|
|
super().__init__()
|
|
self.client = mlflow.tracking.MlflowClient(tracking_uri)
|
|
|
|
experiment = self.client.get_experiment_by_name(experiment_name)
|
|
if experiment is None:
|
|
logger.warning(
|
|
f"Experiment with name f{experiment_name} not found. Creating it."
|
|
)
|
|
self.client.create_experiment(experiment_name)
|
|
experiment = self.client.get_experiment_by_name(experiment_name)
|
|
|
|
run = self.client.create_run(experiment.experiment_id)
|
|
self.run_id = run.info.run_id
|
|
|
|
@rank_zero_only
|
|
def log_hyperparams(self, params):
|
|
for k, v in vars(params).items():
|
|
self.client.log_param(self.run_id, k, v)
|
|
|
|
@rank_zero_only
|
|
def log_metrics(self, metrics, step_num=None):
|
|
timestamp_ms = int(time() * 1000)
|
|
for k, v in metrics.items():
|
|
self.client.log_metric(self.run_id, k, v, timestamp_ms, step_num)
|
|
|
|
def save(self):
|
|
pass
|
|
|
|
@rank_zero_only
|
|
def finalize(self, status="FINISHED"):
|
|
self.client.set_terminated(self.run_id, status)
|