neural_compressor.compression.hpo.search_space

Module Contents

Classes

SearchSpace

Abstract class, the main entrance for search space.

BaseSearchSpace

Base class for Search Space.

DiscreteSearchSpace

Discrete Search Space.

ContinuousSearchSpace

Continuous Search Space.

Functions

register_searchspace(name)

Class decorator to register a SearchSpace subclass to the registry.
class neural_compressor.compression.hpo.search_space.SearchSpace[source]

Abstract class, the main entrance for search space. Using factory pattern to get the actual used class.

Parameters:

  • bound – A tuple or list that limit the max and min number of search space.

  • interval – Only for discrete search space. Intervals in discrete space.

  • value – Only for discrete search space. A list that store all the number for search.

  • type – discrete or continues.

Example

from neural_compressor.compression.hpo import SearchSpace

search_space = {

‘learning_rate’: SearchSpace((0.0001, 0.001)), ‘num_train_epochs’: SearchSpace(bound=(20, 100), interval=1), ‘weight_decay’: SearchSpace((0.0001, 0.001), type=’continuous’) }

neural_compressor.compression.hpo.search_space.register_searchspace(name)[source]

Class decorator to register a SearchSpace subclass to the registry.

Parameters:

  • cls (class) – The subclass of register.

  • name – A string. Define the pruner type.

Returns:

The class of register.

Return type:

cls

class neural_compressor.compression.hpo.search_space.BaseSearchSpace(bound=None, interval=None, value=None, type=None)[source]

Base class for Search Space.

class neural_compressor.compression.hpo.search_space.DiscreteSearchSpace(bound=None, interval=None, value=None, type=None)[source]

Discrete Search Space.

class neural_compressor.compression.hpo.search_space.ContinuousSearchSpace(bound, interval=None, value=None, type=None)[source]

Continuous Search Space.