Upgrade Optuna from v2.x.x to v3.0.0 (#1942)

* Use new distributions in Optuna v3 * Update optuna to v3
kubeflow · Aug 29, 2022 · 733d982 · 733d982
1 parent ca903e2
commit 733d982
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/cmd/suggestion/optuna/v1beta1/requirements.txt b/cmd/suggestion/optuna/v1beta1/requirements.txt
@@ -1,4 +1,4 @@
 grpcio==1.41.1
 protobuf==3.19.1
 googleapis-common-protos==1.53.0
-optuna<3.0.0
+optuna>=3.0.0
diff --git a/pkg/suggestion/v1beta1/optuna/base_service.py b/pkg/suggestion/v1beta1/optuna/base_service.py
@@ -91,9 +91,9 @@ def _get_optuna_search_space(self):
         search_space = {}
         for param in self.search_space.params:
             if param.type == INTEGER:
-                search_space[param.name] = optuna.distributions.IntUniformDistribution(int(param.min), int(param.max))
+                search_space[param.name] = optuna.distributions.IntDistribution(int(param.min), int(param.max))
             elif param.type == DOUBLE:
-                search_space[param.name] = optuna.distributions.UniformDistribution(float(param.min), float(param.max))
+                search_space[param.name] = optuna.distributions.FloatDistribution(float(param.min), float(param.max))
             elif param.type == CATEGORICAL or param.type == DISCRETE:
                 search_space[param.name] = optuna.distributions.CategoricalDistribution(param.list)
         return search_space