Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Apply formatter and follow the Optuna conventions #133

Merged
merged 4 commits into from
Jul 1, 2024
Merged

Apply formatter and follow the Optuna conventions #133

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
00d2002
Apply formatter and follow the Optuna conventions
nabenabe0928 Jun 28, 2024
114edb2
Fix test
nabenabe0928 Jun 28, 2024
35b2190
Apply formatter
nabenabe0928 Jun 28, 2024
bbc0b07
Add future annotations to test
nabenabe0928 Jun 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 81 additions & 58 deletions optuna_integration/comet/comet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,12 +1,9 @@
from __future__ import annotations

from collections.abc import Callable
from collections.abc import Sequence
import functools
import json
from typing import Any
from typing import Callable
from typing import Dict
from typing import Optional
from typing import Sequence
from typing import TYPE_CHECKING
from typing import Union

import optuna
from optuna.study.study import ObjectiveFuncType
Expand All @@ -20,61 +17,87 @@

class CometCallback:
"""
A callback for logging Optuna study trials to a Comet ML Experiment. Comet ML must be installed to run.

This callback is intended for use with Optuna's study.optimize() method. It ensures that all trials
from an Optuna study are logged to a single Comet Experiment, facilitating organized tracking
of hyperparameter optimization. The callback supports both single and multi-objective optimization.

In a distributed training context, where trials from the same study might occur on different machines,
this callback ensures consistency by logging to the same Comet Experiment using an experiment key
stored within the study's user attributes.

By default, Trials are logged as Comet Experiments, which will automatically log code, system metrics,
and many other values. However, it also adds some computational overhead (potentially a few seconds).

Parameters:
- study (optuna.study.Study): The Optuna study object to which the callback is attached.
- workspace (str) - Optional: The workspace in Comet ML where the project resides.
- project_name (str) - Optional: The name of the project in Comet ML where the experiment will be logged. Defaults to "general"
- metric_names ([str]) - Optional: A list of the names of your objective metrics.

Usage:
```
study = optuna.create_study(directions=["minimize", "maximize"])
comet_callback = CometCallback(study, metric_names=["accuracy", "top_k_accuracy"],
project_name="your_project_name", workspace="your_workspace")
study.optimize(your_objective_function, n_trials=100, callbacks=[comet_callback])
```

Note:
The callback checks for an existing Comet Experiment key in the study's user attributes. If present, it initializes
an ExistingExperiment; otherwise, it creates a new APIExperiment and stores its key in the study for future reference.

You will need a Comet API key to log data to Comet.

You can also log extra data directly to your Trial's Experiment via the objective function by using
the @CometCallback.track_in_comet decorator, which exposes an `experiment` property on your trial, like so:
```
study = optuna.create_study(directions=["minimize", "maximize"])
comet_callback = CometCallback(study, metric_names=["accuracy", "top_k_accuracy"],
project_name="your_project_name", workspace="your_workspace")

@comet_callback.track_in_comet()
def your_objective(trial):
trial.experiment.log_other("foo", "bar")
# Rest of your objective function...

study.optimize(your_objective, n_trials=100, callbacks=[comet_callback])
A callback for logging Optuna study trials to a Comet ML Experiment.
Comet ML must be installed to run.

This callback is intended for use with Optuna's study.optimize() method. It ensures that
all trials from an Optuna study are logged to a single Comet Experiment, facilitating organized
tracking of hyperparameter optimization.
The callback supports both single and multi-objective optimization.

In a distributed training context, where trials from the same study might occur on different
machines, this callback ensures consistency by logging to the same Comet Experiment using
an experiment key stored within the study's user attributes.

By default, Trials are logged as Comet Experiments, which will automatically log code,
system metrics, and many other values.
However, it also adds some computational overhead (potentially a few seconds).

Args:
study:
The Optuna study object to which the callback is attached.
workspace:
The workspace in Comet ML where the project resides.
project_name:
The name of the project in Comet ML where the experiment will be logged.
Defaults to ``general``.
metric_names:
A list of the names of your objective metrics.

Example:

Here is an example.

.. code::

study = optuna.create_study(directions=["maximize", "maximize"])
comet_callback = CometCallback(
study,
metric_names=["accuracy", "top_k_accuracy"],
project_name="your_project_name",
workspace="your_workspace",
)
study.optimize(your_objective_function, n_trials=100, callbacks=[comet_callback])

.. note:
The callback checks for an existing Comet Experiment key in the study's user attributes.
If present, it initializes an ExistingExperiment; otherwise,
it creates a new APIExperiment and stores its key in the study for future reference.

You will need a Comet API key to log data to Comet.

You can also log extra data directly to your Trial's Experiment via the objective function
by using the ``@CometCallback.track_in_comet`` decorator,
which exposes an ``experiment`` property on your trial, like so:

.. code::

study = optuna.create_study(directions=["maximize", "maximize"])
comet_callback = CometCallback(
study,
metric_names=["accuracy", "top_k_accuracy"],
project_name="your_project_name",
workspace="your_workspace",
)


@comet_callback.track_in_comet()
def your_objective(trial):
trial.experiment.log_other("foo", "bar")
# Rest of your objective function...


study.optimize(your_objective, n_trials=100, callbacks=[comet_callback])

"""

def __init__(
self,
study: optuna.study.Study,
workspace: Optional[str] = None,
project_name: Optional[str] = "general",
metric_names: Optional[Sequence[str]] = None,
) -> None:
workspace: str | None = None,
project_name: str | None = "general",
metric_names: Sequence[str] | None = None,
):
self._project_name = project_name
self._workspace = workspace
self._study = study
Expand Down Expand Up @@ -195,7 +218,7 @@ def _init_optuna_trial_experiment(
def track_in_comet(self) -> Callable:
def decorator(func: ObjectiveFuncType) -> ObjectiveFuncType:
@functools.wraps(func)
def wrapper(trial: optuna.trial.Trial) -> Union[float, Sequence[float]]:
def wrapper(trial: optuna.trial.Trial) -> float | Sequence[float]:
experiment = self._init_optuna_trial_experiment(self._study, trial)

# Add the experiment to the trial object for easier access for the end-users
Expand Down
27 changes: 5 additions & 22 deletions tests/comet/test_comet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,38 +1,23 @@
from __future__ import annotations

import json
import random
import typing
from typing import Tuple
from unittest import mock

import optuna
import pytest

from optuna_integration.comet import CometCallback


def _objective_func(trial: optuna.trial.Trial) -> float:
x = trial.suggest_float("x", -10, 10)
y = trial.suggest_float("y", 1, 10, log=True)

params = {
"min_samples_leaf": trial.suggest_int("min_samples_leaf", 2, 10),
"max_depth": trial.suggest_int("max_depth", 5, 20),
"min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
}

return (x - 2) ** 2 + (y - 25) ** 2


def _multiobjective_func(trial: optuna.trial.Trial) -> typing.Tuple[float, float]:
def _multiobjective_func(trial: optuna.trial.Trial) -> tuple[float, float]:
x = trial.suggest_float("x", -10, 10)
y = trial.suggest_float("y", 1, 10, log=True)

params = {
"min_samples_leaf": trial.suggest_int("min_samples_leaf", 2, 10),
"max_depth": trial.suggest_int("max_depth", 5, 20),
"min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
}

first_objective = (x - 2) ** 2 + (y - 25) ** 2
second_objective = (x - 2) ** 3 + (y - 25) ** 3
return first_objective, second_objective
Expand Down Expand Up @@ -77,9 +62,7 @@ def test_comet_callback_experiment_key_reuse(
) -> None:
study = optuna.create_study(direction="minimize")
study.set_user_attr("comet_study_experiment_key", "existing_experiment_key")
comet_callback = CometCallback(
study, project_name="optuna_test_reuse", workspace="workspace_reuse"
)
_ = CometCallback(study, project_name="optuna_test_reuse", workspace="workspace_reuse")

# Simulate optimization to check if the existing experiment key is reused
optuna_trial = optuna.trial.create_trial(
Expand Down Expand Up @@ -107,7 +90,7 @@ def test_comet_callback_track_in_comet_decorator(
)

@comet_callback.track_in_comet()
def your_objective(trial: mock.MagicMock) -> Tuple[float, float]:
def your_objective(trial: mock.MagicMock) -> tuple[float, float]:
x = random.uniform(-5, 5)
y = random.uniform(-5, 5)
trial.experiment.log_other("extra_info", "test")
Expand Down
Loading