# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
from typing import Dict, List, Optional, Union
import torch
from overrides import overrides
from archai.discrete_search.api.archai_model import ArchaiModel
from archai.discrete_search.api.model_evaluator import ModelEvaluator
from archai.discrete_search.evaluators.pt_profiler_utils.pt_profiler_eval import profile
[docs]class TorchNumParameters(ModelEvaluator):
"""Total number of parameters."""
def __init__(
self, exclude_cls: Optional[List[torch.nn.Module]] = None, trainable_only: Optional[bool] = True
) -> None:
"""Initialize the evaluator.
Args:
exclude_cls: List of PyTorch module classes to exclude from parameter counting.
trainable_only: A flag indicating whether only trainable parameters should be counted.
"""
self.exclude_cls = exclude_cls
self.trainable_only = trainable_only
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None) -> float:
total_params = sum(
param.numel() for param in model.arch.parameters() if not self.trainable_only or param.requires_grad
)
exclude_params = (
0
if self.exclude_cls is None
else sum(
sum(param.numel() for param in module.parameters())
for module in model.arch.modules()
if isinstance(module, tuple(self.exclude_cls))
)
)
return total_params - exclude_params
[docs]class TorchFlops(ModelEvaluator):
"""Total number of FLOPs."""
def __init__(
self,
forward_args: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
forward_kwargs: Optional[Dict[str, torch.Tensor]] = None,
ignore_layers: Optional[List[str]] = None,
) -> None:
"""Initialize the evaluator.
Args:
forward_args: `model.forward()` arguments used for profilling.
forward_kwargs: `model.forward()` keyword arguments used for profilling.
ignore_layers: List of layer names that should be ignored during profiling.
"""
self.forward_args = forward_args
self.forward_kwargs = forward_kwargs
self.ignore_layers = ignore_layers
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None) -> float:
return profile(
model.arch,
self.forward_args,
self.forward_kwargs,
num_warmups=0,
num_samples=1,
ignore_layers=self.ignore_layers,
)["flops"]
[docs]class TorchMacs(ModelEvaluator):
"""Total number of MACs."""
def __init__(
self,
forward_args: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
forward_kwargs: Optional[Dict[str, torch.Tensor]] = None,
ignore_layers: Optional[List[str]] = None,
) -> None:
"""Initialize the evaluator.
Args:
forward_args: `model.forward()` arguments used for profilling.
forward_kwargs: `model.forward()` keyword arguments used for profilling.
ignore_layers: List of layer names that should be ignored during profiling.
"""
self.forward_args = forward_args
self.forward_kwargs = forward_kwargs
self.ignore_layers = ignore_layers
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None) -> float:
return profile(
model.arch,
self.forward_args,
self.forward_kwargs,
num_warmups=0,
num_samples=1,
ignore_layers=self.ignore_layers,
)["macs"]
[docs]class TorchLatency(ModelEvaluator):
"""Average/median latency (in seconds) of a PyTorch model using a sample input."""
def __init__(
self,
forward_args: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
forward_kwargs: Optional[Dict[str, torch.Tensor]] = None,
num_warmups: Optional[int] = 1,
num_samples: Optional[int] = 1,
use_cuda: Optional[bool] = False,
use_median: Optional[bool] = False,
ignore_layers: Optional[List[str]] = None,
) -> None:
"""Initialize the evaluator.
Args:
forward_args: `model.forward()` arguments used for profilling.
forward_kwargs: `model.forward()` keyword arguments used for profilling.
num_warmups: Number of warmup runs before profilling.
num_samples: Number of runs after warmup.
use_cuda: Whether to use CUDA instead of CPU.
use_median: Whether to use median instead of mean to average memory and latency.
ignore_layers: List of layer names that should be ignored during profiling.
"""
self.forward_args = forward_args
self.forward_kwargs = forward_kwargs
self.ignore_layers = ignore_layers
self.num_warmups = num_warmups
self.num_samples = num_samples
self.use_cuda = use_cuda
self.use_median = use_median
self.ignore_layers = ignore_layers
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None) -> float:
return profile(
model.arch,
self.forward_args,
self.forward_kwargs,
num_warmups=self.num_warmups,
num_samples=self.num_samples,
use_cuda=self.use_cuda,
use_median=self.use_median,
ignore_layers=self.ignore_layers,
)["latency"]
[docs]class TorchPeakCudaMemory(ModelEvaluator):
"""Measures CUDA peak memory (in bytes) of a PyTorch model using a sample input."""
def __init__(
self,
forward_args: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
forward_kwargs: Optional[Dict[str, torch.Tensor]] = None,
num_warmups: Optional[int] = 1,
num_samples: Optional[int] = 1,
use_median: Optional[bool] = False,
ignore_layers: Optional[List[str]] = None,
) -> None:
"""Initialize the evaluator.
Args:
forward_args: `model.forward()` arguments used for profilling.
forward_kwargs: `model.forward()` keyword arguments used for profilling.
num_warmups: Number of warmup runs before profilling.
num_samples: Number of runs after warmup.
use_median: Whether to use median instead of mean to average memory and latency.
ignore_layers: List of layer names that should be ignored during profiling.
"""
self.forward_args = forward_args
self.forward_kwargs = forward_kwargs
self.ignore_layers = ignore_layers
self.num_warmups = num_warmups
self.num_samples = num_samples
self.use_median = use_median
self.ignore_layers = ignore_layers
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None) -> float:
return profile(
model.arch,
self.forward_args,
self.forward_kwargs,
num_warmups=self.num_warmups,
num_samples=self.num_samples,
use_cuda=True,
use_median=self.use_median,
ignore_layers=self.ignore_layers,
)["peak_memory"]
[docs]class TorchPeakCpuMemory(ModelEvaluator):
"""Measures CPU peak memory (in bytes) of a PyTorch model using a sample input."""
def __init__(
self,
forward_args: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
forward_kwargs: Optional[Dict[str, torch.Tensor]] = None,
):
"""Initialize the evaluator.
Args:
forward_args: `model.forward()` arguments used for profilling.
forward_kwargs: `model.forward()` keyword arguments used for profilling.
"""
forward_args = forward_args if forward_args is not None else []
self.forward_args = [forward_args] if isinstance(forward_args, torch.Tensor) else forward_args
self.forward_kwargs = forward_kwargs or {}
[docs] @overrides
def evaluate(self, model: ArchaiModel, budget: Optional[float] = None):
model.arch.to("cpu")
forward_args = tuple([arg.to("cpu") for arg in self.forward_args])
forward_kwargs = {key: value.to("cpu") for key, value in self.forward_kwargs.items()}
is_training = model.arch.training
model.arch.eval()
with torch.profiler.profile(
activities=[torch.profiler.ProfilerActivity.CPU], record_shapes=True, profile_memory=True
) as prof:
with torch.profiler.record_function("model_inference"):
model.arch(*forward_args, **forward_kwargs)
event_list = prof.key_averages()
peak_memory = max(event.cpu_memory_usage for event in event_list)
if is_training:
model.arch.train()
return peak_memory
