OptimizationAgent (pro)

The OptimizationAgent is the main class in pruna_pro. It is used to optimize your model.

Usage example OptimizationAgent

from pruna_pro import OptimizationAgent
from pruna.evaluation.task import Task
from pruna.data.pruna_datamodule import PrunaDataModule
from diffusers import StableDiffusionPipeline

# Define your task with metrics and your model
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
task = Task(['elapsed_time'], datamodule=PrunaDataModule.from_string('LAION256'))

# Create the optimization agent with the model and your objectives
agent = OptimizationAgent(model=model, task=task)

# Find the best configuration instantly
optimized_model = agent.instant_search()

Class API OptimizationAgent

class OptimizationAgent(model, task, token=None, device_cast_fn=None)

OptimizationAgent automatically finds the best algorithms and hyperparameters for a given model and objective.

Parameters:

  • model (Any) – The model to be optimized.

  • task (Task) – The task to be used.

  • token (str | None, optional) – The token to be used.

  • device_cast_fn (Callable | None, optional) – The device cast function to be used if the default device cast function fails.

__init__(model, task, token=None, device_cast_fn=None)
Parameters:

  • model (Any)

  • task (Task)

  • token (str | None)

  • device_cast_fn (Callable | None)

Return type:

None

Suggest the best configuration for a given model on CUDA instantly.

Parameters:

  • return_config (bool, optional) – Whether to return the configuration.

  • verbose (bool, optional) – Whether to print verbose output.

Returns:

The optimized model.

Return type:

Any

Find the best configuration for the model using the probabilistic search strategy on CUDA.

Parameters:

  • n_trials (int, optional) – The number of trials to run.

  • n_iterations (int, optional) – The number of iterations to run for the evaluation of the given metrics.

  • n_warmup_iterations (int, optional) – The number of warmup iterations to run for the evaluation of the given metrics.

  • return_config (bool, optional) – Whether to return the configuration.

  • verbose (bool, optional) – Whether to print verbose output.

  • seed (int, optional) – The seed to be used.

Returns:

The optimized model.

Return type:

Any