Skip to content

Commit 853dc71

Browse files
Corrected the docstring of TuningConfig (#1639)
Signed-off-by: yiliu30 <yi4.liu@intel.com> Co-authored-by: Kaihui-intel <kaihui.tang@intel.com>
1 parent 2b86e50 commit 853dc71

File tree

1 file changed

+19
-16
lines changed

1 file changed

+19
-16
lines changed

neural_compressor/common/base_tuning.py

+19-16
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,11 @@
1414

1515

1616
import copy
17-
import inspect
1817
import uuid
1918
from typing import Any, Callable, Dict, Generator, Iterator, List, Optional, Sized, Tuple, Union
2019

2120
from neural_compressor.common import Logger
22-
from neural_compressor.common.base_config import BaseConfig, ComposableConfig
21+
from neural_compressor.common.base_config import BaseConfig
2322
from neural_compressor.common.utils import TuningLogger
2423

2524
logger = Logger().get_logger()
@@ -227,19 +226,11 @@ def __iter__(self) -> Generator[BaseConfig, Any, None]:
227226

228227

229228
class TuningConfig:
230-
"""Base Class for Tuning Criterion.
231-
232-
Args:
233-
config_set: quantization configs. Default value is empty.
234-
A single config or a list of configs. More details can
235-
be found in the `from_fwk_configs`of `ConfigSet` class.
236-
max_trials: Max tuning times. Default value is 100. Combine with timeout field to decide when to exit.
237-
tolerable_loss: This float indicates how much metric loss we can accept. \
238-
The metric loss is relative, it can be both positive and negative. Default is 0.01.
229+
"""Config for auto tuning pipeline.
239230
240231
Examples:
241232
# TODO: to refine it
242-
from neural_compressor import TuningConfig
233+
from neural_compressor.torch.quantization import TuningConfig
243234
tune_config = TuningConfig(
244235
config_set=[config1, config2, ...],
245236
max_trials=3,
@@ -264,13 +255,25 @@ class TuningConfig:
264255
"""
265256

266257
def __init__(
267-
self, config_set=None, max_trials=100, sampler: Sampler = default_sampler, tolerable_loss=0.01
268-
) -> None:
269-
"""Init a TuneCriterion object."""
258+
self,
259+
config_set: Union[BaseConfig, List[BaseConfig]] = None,
260+
sampler: Sampler = default_sampler,
261+
tolerable_loss=0.01,
262+
max_trials=100,
263+
):
264+
"""Initial a TuningConfig.
265+
266+
Args:
267+
config_set: A single config or a list of configs. Defaults to None.
268+
sampler: tuning sampler that decide the trials order. Defaults to default_sampler.
269+
tolerable_loss: This float indicates how much metric loss we can accept.
270+
The metric loss is relative, it can be both positive and negative. Default is 0.01.
271+
max_trials: Max tuning times. Combine with `tolerable_loss` field to decide when to stop. Default is 100.
272+
"""
270273
self.config_set = config_set
271-
self.max_trials = max_trials
272274
self.sampler = sampler
273275
self.tolerable_loss = tolerable_loss
276+
self.max_trials = max_trials
274277

275278

276279
class _TrialRecord:

0 commit comments

Comments
 (0)