MLFlowLogger

class torchtune.training.metric_logging.MLFlowLogger(experiment_name: Optional[str] = None, tracking_uri: Optional[str] = None, run_id: Optional[str] = None, run_name: Optional[str] = None)

Logger for use w/ MLFlow (https://mlflow.org/).

Parameters:

• experiment_name (Optional[str]) – MLFlow experiment name. If not specified, will default to MLFLOW_EXPERIMENT_NAME environment variable if set, or default.

• tracking_uri (Optional[str]) – MLFlow tracking uri. If not specified, will default to MLFLOW_TRACKING_URI environment variable if set, or default.

• run_id (Optional[str]) – MLFlow run name. If not specified, will default to mlflow-generated HRID. Unused if run_id is specified or MLFLOW_RUN_ID environment variable is found.

• run_name (Optional[str]) – MLFlow run ID. If not specified, will default to MLFLOW_RUN_ID environment variable if set, or a new run will be created.

Example

>>> logger = MLFlowLogger(experiment_name="my_experiment", run_name="run1")
>>> logger.log("accuracy", 0.95, step=1)
>>> logger.log_dict({"loss": 0.1, "accuracy": 0.95}, step=1)
>>> logger.log_config(config)
>>> logger.close()

Raises:

ImportError – If mlflow package is not installed.

Note

This logger requires the mlflow package to be installed. You can install it with pip install mlflow.

close() → None

Ends the MLflow run. After calling close, no further logging should be performed.

log(name: str, data: Union[Tensor, ndarray, int, float], step: int) → None

Log scalar data.

Parameters:

• name (str) – tag name used to group scalars

• data (Scalar) – scalar data to log

• step (int) – step value to record

log_config(config: DictConfig) → None

Saves the config locally and also logs the config to mlflow. The config is stored in the same directory as the checkpoint.

Parameters:

config (DictConfig) – config to log

log_dict(payload: Mapping[str, Union[Tensor, ndarray, int, float]], step: int) → None

Log multiple scalar values.

Parameters:

• payload (Mapping[str, Scalar]) – dictionary of tag name and scalar value

• step (int) – step value to record