SmashConfig

SmashConfig is an essential tool in pruna for configuring parameters to optimize your models. The SmashConfig contains configuration for pruna. You can see the configuration for different components in the Customize components section.

Usage examples SmashConfig

This manual explains how to define and use a SmashConfig.

from pruna import SmashConfig
smash_config = SmashConfig()

After creating an empty SmashConfig, you can set activate a algorithm by adding it to the SmashConfig:

smash_config['quantizer'] = 'hqq'

Additionally, you can overwrite the defaults of the Algorithms you have added by setting the hyperparameters in the SmashConfig:

smash_config['hqq_weight_bits'] = 4

You’re done! You created your SmashConfig and can now pass it to the smash function.

Class API SmashConfig

class SmashConfig(configuration=None, batch_size=1, device=None, cache_dir_prefix=PosixPath('/home/docs/.cache/pruna'))

Bases: object

Wrapper class to hold a ConfigSpace Configuration object as a Smash configuration.

Parameters:
  • configuration (list[str] | Dict[str, Any] | Configuration | None, optional) – The configuration to be used for smashing. If None, a default configuration will be created.

  • batch_size (int, optional) – The number of batches to process at once. Default is 1.

  • device (str | torch.device | None, optional) – The device to be used for smashing, options are “cpu”, “cuda”, “mps”, “accelerate”. Default is None. If None, the best available device will be used.

  • cache_dir_prefix (str, optional) – The prefix for the cache directory. If None, a default cache directory will be created.

classmethod from_list(configuration, batch_size=1, device=None, cache_dir_prefix=PosixPath('/home/docs/.cache/pruna'))

Create a SmashConfig from a list of algorithm names.

Parameters:
  • configuration (list[str]) – The list of algorithm names to create the SmashConfig with.

  • batch_size (int, optional) – The batch size to use for the SmashConfig. Default is 1.

  • device (str | torch.device | None, optional) – The device to use for the SmashConfig. Default is None.

  • cache_dir_prefix (str | Path, optional) – The prefix for the cache directory. Default is DEFAULT_CACHE_DIR.

Returns:

The SmashConfig object instantiated from the list.

Return type:

SmashConfig

Examples

>>> config = SmashConfig.from_list(["fastercache", "diffusers_int8"])
>>> config
SmashConfig(
 'fastercache': True,
 'diffusers_int8': True,
)
classmethod from_dict(configuration, batch_size=1, device=None, cache_dir_prefix=PosixPath('/home/docs/.cache/pruna'))

Create a SmashConfig from a dictionary of algorithms and their hyperparameters.

Parameters:
  • configuration (Dict[str, Any]) – The dictionary to create the SmashConfig from.

  • batch_size (int, optional) – The batch size to use for the SmashConfig. Default is 1.

  • device (str | torch.device | None, optional) – The device to use for the SmashConfig. Default is None.

  • cache_dir_prefix (str | Path, optional) – The prefix for the cache directory. Default is DEFAULT_CACHE_DIR.

Returns:

The SmashConfig object instantiated from the dictionary.

Return type:

SmashConfig

Examples

>>> config = SmashConfig.from_dict({"fastercache": True, "diffusers_int8": True})
>>> config
SmashConfig(
 'fastercache': True,
 'diffusers_int8': True,
)
cleanup_cache_dir()

Clean up the cache directory.

Return type:

None

reset_cache_dir()

Reset the cache directory.

Return type:

None

load_from_json(path)

Load a SmashConfig from a JSON file.

Parameters:

path (str| Path) – The file path to the JSON file containing the configuration.

Return type:

None

save_to_json(path)

Save the SmashConfig to a JSON file, including additional keys.

Parameters:

path (str| Path]) – The file path where the JSON file will be saved.

Return type:

None

flush_configuration()

Remove all algorithm hyperparameters from the SmashConfig.

Examples

>>> config = SmashConfig(["fastercache", "diffusers_int8"])
>>> config.flush_configuration()
>>> config
SmashConfig()
Return type:

None

train_dataloader(**kwargs)

Getter for the train DataLoader instance.

Parameters:

**kwargs (dict) – Any additional arguments used when loading data, overriding the default values provided in the constructor. Examples: img_size: int would override the collate function default for image generation, while batch_size: int, shuffle: bool, pin_memory: bool, … would override the dataloader defaults.

Returns:

The DataLoader instance associated with the SmashConfig.

Return type:

torch.utils.data.DataLoader | None

val_dataloader(**kwargs)

Getter for the validation DataLoader instance.

Parameters:

**kwargs (dict) – Any additional arguments used when loading data, overriding the default values provided in the constructor. Examples: img_size: int would override the collate function default for image generation, while batch_size: int, shuffle: bool, pin_memory: bool, … would override the dataloader defaults.

Returns:

The DataLoader instance associated with the SmashConfig.

Return type:

torch.utils.data.DataLoader | None

test_dataloader(**kwargs)

Getter for the test DataLoader instance.

Parameters:

**kwargs (dict) – Any additional arguments used when loading data, overriding the default values provided in the constructor. Examples: img_size: int would override the collate function default for image generation, while batch_size: int, shuffle: bool, pin_memory: bool, … would override the dataloader defaults.

Returns:

The DataLoader instance associated with the SmashConfig.

Return type:

torch.utils.data.DataLoader | None

add_data(arg)
add_data(dataset_name, *args, **kwargs)
add_data(datasets, collate_fn, *args, **kwargs)
add_data(datasets, collate_fn, *args, **kwargs)
add_data(datamodule)

Add data to the SmashConfig.

Parameters:

arg (Any) – The argument to be used.

add_tokenizer(tokenizer)

Add a tokenizer to the SmashConfig.

Parameters:

tokenizer (str | transformers.AutoTokenizer) – The tokenizer to be added to the SmashConfig.

Return type:

None

add_processor(processor)

Add a processor to the SmashConfig.

Parameters:

processor (str | transformers.AutoProcessor) – The processor to be added to the SmashConfig.

Return type:

None

add_target_module(target_module)

Add a target module to prune to the SmashConfig.

Parameters:

target_module (Any) – The target module to prune.

Return type:

None

lock_batch_size()

Lock the batch size in the SmashConfig.

Return type:

None

is_batch_size_locked()

Check if the batch size is locked in the SmashConfig.

Returns:

True if the batch size is locked, False otherwise.

Return type:

bool

add(request)

Add an algorithm or specify the hyperparameters of an algorithm to the SmashConfig.

Parameters:

request (str | list[str] | dict[str, Any]) – The value to add to the SmashConfig.

Return type:

None

Examples

>>> config = SmashConfig()
>>> config = SmashConfig()
>>> config.add("fastercache")
>>> config.add("diffusers_int8")
>>> config
SmashConfig(
 'fastercache': True,
 'diffusers_int8': True,
)
get_active_algorithms()

Get all active algorithms in this smash config.

Returns:

The active algorithms in this smash config.

Return type:

list[str]

overwrite_algorithm_order(algorithm_order)

Overwrite the graph-induced order of algorithms if desired.

Parameters:

algorithm_order (list[str]) – The order of algorithms to be applied.

Return type:

None

disable_saving()

Disable the saving of the SmashConfig to speed up the smashing process.

Return type:

None