profiler

shrinkai.profiler

Profiling: measuring and comparing a model's deployment footprint.

Profiler.compare (returning a BenchmarkReport) is the main entry point, gathering parameter count, disk size, latency/throughput, optional accuracy, and optional FLOPs (count_flops) for a teacher/student pair in one call, also reachable as Distiller.benchmark, Pruner.benchmark, ChannelPruner.benchmark, and Quantizer.benchmark. The individual measuring functions (measure_latency, count_parameters, estimate_model_size_mb, get_process_ram_mb, get_device_memory_mb, count_flops) are also usable standalone.

Modules:

Name	Description
`accuracy`
`benchmark`
`flops`	FLOPs (floating point operations) counting for edge-deployment profiling.
`latency`
`memory`

Classes:

Name	Description
`BenchmarkReport`	Holds comparison results between Teacher and Student models.
`ModelProfile`	Container holding profiled metrics for an individual model.
`Profiler`	Benchmark runner comparing Teacher and Student architectures.

Functions:

Name	Description
`count_flops`	Counts the total FLOPs of one forward pass of `model` on `sample_input`.
`count_parameters`	Counts total, trainable, and non-trainable parameters.
`estimate_model_size_mb`	Estimates serialized state dictionary size on disk in Megabytes (MB).
`get_device_memory_mb`	Returns the currently allocated memory on the specified hardware accelerator.
`get_process_ram_mb`	Returns current host RAM consumption of the running Python process.
`measure_latency`	Measures average inference latency and throughput (FPS / samples per second).

Classes

BenchmarkReport

Holds comparison results between Teacher and Student models.

Methods:

Name	Description
`show`	Renders an interactive formatted comparison table in the console.

Source code in src/shrinkai/profiler/benchmark.py

class BenchmarkReport:
    """Holds comparison results between Teacher and Student models."""

    def __init__(self, teacher_profile: ModelProfile, student_profile: ModelProfile) -> None:
        self.teacher = teacher_profile
        self.student = student_profile

    def _compute_gains(self) -> dict[str, str]:
        """Calculates percentage and factor improvements."""
        if self.teacher.total_params > 0:
            param_reduction = (1 - self.student.total_params / self.teacher.total_params) * 100
        else:
            param_reduction = 0.0

        if self.teacher.size_mb > 0:
            size_reduction = (1 - self.student.size_mb / self.teacher.size_mb) * 100
        else:
            size_reduction = 0.0

        size_multiplier = (
            self.teacher.size_mb / self.student.size_mb if self.student.size_mb > 0 else 1.0
        )
        speedup = self.student.fps / self.teacher.fps if self.teacher.fps > 0 else 1.0

        return {
            "param_reduction": f"-{param_reduction:.1f}%",
            "size_reduction": f"-{size_reduction:.1f}% ({size_multiplier:.1f}x smaller)",
            "speedup": f"+{speedup:.1f}x ({self.student.fps:.1f} FPS)",
        }

    def show(self) -> None:
        """Renders an interactive formatted comparison table in the console."""
        console = Console()
        gains = self._compute_gains()

        table = Table(title="Benchmark Report", header_style="bold cyan")
        table.add_column("Metric", style="bold")
        table.add_column(f"Teacher ({self.teacher.name})", justify="right")
        table.add_column(f"Student ({self.student.name})", justify="right")
        table.add_column("Gain / Compression", justify="right", style="green")

        table.add_row(
            "Parameters",
            f"{self.teacher.total_params / 1e6:.2f} M",
            f"{self.student.total_params / 1e6:.2f} M",
            gains["param_reduction"],
        )
        table.add_row(
            "Model Size (Disk)",
            f"{self.teacher.size_mb:.2f} MB",
            f"{self.student.size_mb:.2f} MB",
            gains["size_reduction"],
        )
        latency_speedup = (
            self.teacher.sample_latency_ms / self.student.sample_latency_ms
            if self.student.sample_latency_ms > 0
            else 1.0
        )
        table.add_row(
            "Latency / Sample",
            f"{self.teacher.sample_latency_ms:.2f} ms",
            f"{self.student.sample_latency_ms:.2f} ms",
            f"{latency_speedup:.1f}x faster",
        )
        table.add_row(
            "Throughput (FPS)",
            f"{self.teacher.fps:.1f} img/s",
            f"{self.student.fps:.1f} img/s",
            gains["speedup"],
        )

        if self.teacher.flops is not None and self.student.flops is not None:
            flops_reduction = (
                (1 - self.student.flops / self.teacher.flops) * 100
                if self.teacher.flops > 0
                else 0.0
            )
            table.add_row(
                "FLOPs / Sample",
                f"{self.teacher.flops / 1e6:.2f} MFLOPs",
                f"{self.student.flops / 1e6:.2f} MFLOPs",
                f"-{flops_reduction:.1f}%",
            )

        if self.teacher.accuracy is not None and self.student.accuracy is not None:
            retention = (
                (self.student.accuracy / self.teacher.accuracy) * 100
                if self.teacher.accuracy > 0
                else 0.0
            )
            table.add_row(
                "Accuracy",
                f"{self.teacher.accuracy:.2f}%",
                f"{self.student.accuracy:.2f}%",
                f"{retention:.1f}% retained",
            )

        console.print(table)

Methods:

show

show() -> None

Renders an interactive formatted comparison table in the console.

Source code in src/shrinkai/profiler/benchmark.py

def show(self) -> None:
    """Renders an interactive formatted comparison table in the console."""
    console = Console()
    gains = self._compute_gains()

    table = Table(title="Benchmark Report", header_style="bold cyan")
    table.add_column("Metric", style="bold")
    table.add_column(f"Teacher ({self.teacher.name})", justify="right")
    table.add_column(f"Student ({self.student.name})", justify="right")
    table.add_column("Gain / Compression", justify="right", style="green")

    table.add_row(
        "Parameters",
        f"{self.teacher.total_params / 1e6:.2f} M",
        f"{self.student.total_params / 1e6:.2f} M",
        gains["param_reduction"],
    )
    table.add_row(
        "Model Size (Disk)",
        f"{self.teacher.size_mb:.2f} MB",
        f"{self.student.size_mb:.2f} MB",
        gains["size_reduction"],
    )
    latency_speedup = (
        self.teacher.sample_latency_ms / self.student.sample_latency_ms
        if self.student.sample_latency_ms > 0
        else 1.0
    )
    table.add_row(
        "Latency / Sample",
        f"{self.teacher.sample_latency_ms:.2f} ms",
        f"{self.student.sample_latency_ms:.2f} ms",
        f"{latency_speedup:.1f}x faster",
    )
    table.add_row(
        "Throughput (FPS)",
        f"{self.teacher.fps:.1f} img/s",
        f"{self.student.fps:.1f} img/s",
        gains["speedup"],
    )

    if self.teacher.flops is not None and self.student.flops is not None:
        flops_reduction = (
            (1 - self.student.flops / self.teacher.flops) * 100
            if self.teacher.flops > 0
            else 0.0
        )
        table.add_row(
            "FLOPs / Sample",
            f"{self.teacher.flops / 1e6:.2f} MFLOPs",
            f"{self.student.flops / 1e6:.2f} MFLOPs",
            f"-{flops_reduction:.1f}%",
        )

    if self.teacher.accuracy is not None and self.student.accuracy is not None:
        retention = (
            (self.student.accuracy / self.teacher.accuracy) * 100
            if self.teacher.accuracy > 0
            else 0.0
        )
        table.add_row(
            "Accuracy",
            f"{self.teacher.accuracy:.2f}%",
            f"{self.student.accuracy:.2f}%",
            f"{retention:.1f}% retained",
        )

    console.print(table)

ModelProfile `dataclass`

Container holding profiled metrics for an individual model.

Source code in src/shrinkai/profiler/benchmark.py

@dataclass
class ModelProfile:
    """Container holding profiled metrics for an individual model."""

    name: str
    total_params: int
    size_mb: float
    sample_latency_ms: float
    fps: float
    accuracy: float | None = None
    flops: int | None = None

Profiler

Benchmark runner comparing Teacher and Student architectures.

Methods:

Name	Description
`compare`	Executes complete profiling suite on both models and generates comparison.

Source code in src/shrinkai/profiler/benchmark.py

class Profiler:
    """Benchmark runner comparing Teacher and Student architectures."""

    @staticmethod
    def compare(
        teacher: nn.Module,
        student: nn.Module,
        sample_input: torch.Tensor,
        device: torch.device | str = "auto",
        teacher_name: str = "Teacher",
        student_name: str = "Student",
        teacher_acc: float | None = None,
        student_acc: float | None = None,
        compute_flops: bool = False,
    ) -> BenchmarkReport:
        """Executes complete profiling suite on both models and generates comparison.

        Args:
            teacher: Teacher PyTorch model.
            student: Student PyTorch model.
            sample_input: Representative tensor input batch.
            device: Device target ('auto', 'mps', 'cuda', 'cpu').
            teacher_name: Display label for teacher.
            student_name: Display label for student.
            teacher_acc: Optional pre-computed teacher accuracy.
            student_acc: Optional pre-computed student accuracy.
            compute_flops: If True, also counts and reports FLOPs per sample for
                both models (see `count_flops`). Defaults to False, since custom/
                opaque ops (e.g. quantized kernels) are silently undercounted as
                0 FLOPs, opt in only when both models use standard ops.

        Returns:
            BenchmarkReport: Structured report ready for `.show()`.
        """
        resolved_device = resolve_device(device)

        t_params = count_parameters(teacher)["total_params"]
        t_size = estimate_model_size_mb(teacher)
        t_lat = measure_latency(teacher, sample_input, resolved_device)
        t_flops = count_flops(teacher, sample_input, resolved_device) if compute_flops else None
        teacher_profile = ModelProfile(
            name=teacher_name,
            total_params=t_params,
            size_mb=t_size,
            sample_latency_ms=t_lat["sample_latency_ms"],
            fps=t_lat["fps"],
            accuracy=teacher_acc,
            flops=t_flops,
        )

        s_params = count_parameters(student)["total_params"]
        s_size = estimate_model_size_mb(student)
        s_lat = measure_latency(student, sample_input, resolved_device)
        s_flops = count_flops(student, sample_input, resolved_device) if compute_flops else None
        student_profile = ModelProfile(
            name=student_name,
            total_params=s_params,
            size_mb=s_size,
            flops=s_flops,
            sample_latency_ms=s_lat["sample_latency_ms"],
            fps=s_lat["fps"],
            accuracy=student_acc,
        )

        return BenchmarkReport(teacher_profile, student_profile)

Methods:

compare `staticmethod`

compare(
    teacher: Module,
    student: Module,
    sample_input: Tensor,
    device: device | str = "auto",
    teacher_name: str = "Teacher",
    student_name: str = "Student",
    teacher_acc: float | None = None,
    student_acc: float | None = None,
    compute_flops: bool = False,
) -> BenchmarkReport

Executes complete profiling suite on both models and generates comparison.

Parameters:

Name	Type	Description	Default
`teacher`	`Module`	Teacher PyTorch model.	required
`student`	`Module`	Student PyTorch model.	required
`sample_input`	`Tensor`	Representative tensor input batch.	required
`device`	`device \| str`	Device target ('auto', 'mps', 'cuda', 'cpu').	`'auto'`
`teacher_name`	`str`	Display label for teacher.	`'Teacher'`
`student_name`	`str`	Display label for student.	`'Student'`
`teacher_acc`	`float \| None`	Optional pre-computed teacher accuracy.	`None`
`student_acc`	`float \| None`	Optional pre-computed student accuracy.	`None`
`compute_flops`	`bool`	If True, also counts and reports FLOPs per sample for both models (see `count_flops`). Defaults to False, since custom/ opaque ops (e.g. quantized kernels) are silently undercounted as 0 FLOPs, opt in only when both models use standard ops.	`False`

Returns:

Name	Type	Description
`BenchmarkReport`	`BenchmarkReport`	Structured report ready for `.show()`.

Source code in src/shrinkai/profiler/benchmark.py

@staticmethod
def compare(
    teacher: nn.Module,
    student: nn.Module,
    sample_input: torch.Tensor,
    device: torch.device | str = "auto",
    teacher_name: str = "Teacher",
    student_name: str = "Student",
    teacher_acc: float | None = None,
    student_acc: float | None = None,
    compute_flops: bool = False,
) -> BenchmarkReport:
    """Executes complete profiling suite on both models and generates comparison.

    Args:
        teacher: Teacher PyTorch model.
        student: Student PyTorch model.
        sample_input: Representative tensor input batch.
        device: Device target ('auto', 'mps', 'cuda', 'cpu').
        teacher_name: Display label for teacher.
        student_name: Display label for student.
        teacher_acc: Optional pre-computed teacher accuracy.
        student_acc: Optional pre-computed student accuracy.
        compute_flops: If True, also counts and reports FLOPs per sample for
            both models (see `count_flops`). Defaults to False, since custom/
            opaque ops (e.g. quantized kernels) are silently undercounted as
            0 FLOPs, opt in only when both models use standard ops.

    Returns:
        BenchmarkReport: Structured report ready for `.show()`.
    """
    resolved_device = resolve_device(device)

    t_params = count_parameters(teacher)["total_params"]
    t_size = estimate_model_size_mb(teacher)
    t_lat = measure_latency(teacher, sample_input, resolved_device)
    t_flops = count_flops(teacher, sample_input, resolved_device) if compute_flops else None
    teacher_profile = ModelProfile(
        name=teacher_name,
        total_params=t_params,
        size_mb=t_size,
        sample_latency_ms=t_lat["sample_latency_ms"],
        fps=t_lat["fps"],
        accuracy=teacher_acc,
        flops=t_flops,
    )

    s_params = count_parameters(student)["total_params"]
    s_size = estimate_model_size_mb(student)
    s_lat = measure_latency(student, sample_input, resolved_device)
    s_flops = count_flops(student, sample_input, resolved_device) if compute_flops else None
    student_profile = ModelProfile(
        name=student_name,
        total_params=s_params,
        size_mb=s_size,
        flops=s_flops,
        sample_latency_ms=s_lat["sample_latency_ms"],
        fps=s_lat["fps"],
        accuracy=student_acc,
    )

    return BenchmarkReport(teacher_profile, student_profile)

Functions:

count_flops

count_flops(
    model: Module,
    sample_input: Tensor | tuple[Tensor, ...],
    device: device | str = "auto",
) -> int

Counts the total FLOPs of one forward pass of model on sample_input.

Uses torch.utils.flop_counter.FlopCounterMode, which instruments the actual tensor operations dispatched during the forward pass, covering standard layers (Conv, Linear, matmul, attention, ...) precisely, rather than a hand-maintained per-layer-type formula. 1 multiply-add (MAC) is counted as 2 FLOPs, matching the usual convention.

Note

Custom/opaque kernels (e.g. the quantized ops produced by shrinkai.compression.quantization.Quantizer) have no FLOPs formula registered and are silently counted as 0. A warning is logged if the total comes back as 0 despite the model having parameters, since that usually signals an undercount rather than a genuinely free model.

Parameters:

Name	Type	Description	Default
`model`	`Module`	Model to profile.	required
`sample_input`	`Tensor \| tuple[Tensor, ...]`	Representative input tensor (or tuple of tensors, for multi-input models) matching the model's forward signature.	required
`device`	`device \| str`	Device to run the single, untimed forward pass on. FLOPs are a static property of the computation graph and do not depend on the device; this only needs to be a device the model can actually run on.	`'auto'`

Returns:

Name	Type	Description
`int`	`int`	Total FLOPs for one forward pass on `sample_input`.

Source code in src/shrinkai/profiler/flops.py

def count_flops(
    model: nn.Module,
    sample_input: torch.Tensor | tuple[torch.Tensor, ...],
    device: torch.device | str = "auto",
) -> int:
    """Counts the total FLOPs of one forward pass of `model` on `sample_input`.

    Uses `torch.utils.flop_counter.FlopCounterMode`, which instruments the actual
    tensor operations dispatched during the forward pass, covering standard
    layers (Conv, Linear, matmul, attention, ...) precisely, rather than a
    hand-maintained per-layer-type formula. 1 multiply-add (MAC) is counted as 2
    FLOPs, matching the usual convention.

    Note:
        Custom/opaque kernels (e.g. the quantized ops produced by
        `shrinkai.compression.quantization.Quantizer`) have no FLOPs formula
        registered and are silently counted as 0. A warning is logged if the
        total comes back as 0 despite the model having parameters, since that
        usually signals an undercount rather than a genuinely free model.

    Args:
        model: Model to profile.
        sample_input: Representative input tensor (or tuple of tensors, for
            multi-input models) matching the model's forward signature.
        device: Device to run the single, untimed forward pass on. FLOPs are a
            static property of the computation graph and do not depend on the
            device; this only needs to be a device the model can actually run on.

    Returns:
        int: Total FLOPs for one forward pass on `sample_input`.
    """
    resolved_device = resolve_device(device)
    model = model.to(resolved_device)
    args = sample_input if isinstance(sample_input, tuple) else (sample_input,)
    args = tuple(a.to(resolved_device) if isinstance(a, torch.Tensor) else a for a in args)

    was_training = model.training
    model.eval()

    with FlopCounterMode(display=False) as flop_counter:
        with torch.no_grad():
            model(*args)

    model.train(was_training)

    total_flops = flop_counter.get_total_flops()
    if total_flops == 0 and any(p.numel() > 0 for p in model.parameters()):
        logger.warning(
            "count_flops returned 0 for a model with parameters. This usually "
            "means it uses custom/opaque ops (e.g. quantized kernels) that "
            "FlopCounterMode cannot see through, not that it is actually free."
        )

    return total_flops

count_parameters

count_parameters(model: Module) -> dict[str, int]

Counts total, trainable, and non-trainable parameters.

Parameters:

Name	Type	Description	Default
`model`	`Module`	PyTorch model.	required

Returns:

Type	Description
`dict[str, int]`	dict[str, int]: Parameter count breakdown.

Source code in src/shrinkai/profiler/memory.py

def count_parameters(model: nn.Module) -> dict[str, int]:
    """Counts total, trainable, and non-trainable parameters.

    Args:
        model: PyTorch model.

    Returns:
        dict[str, int]: Parameter count breakdown.
    """
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {
        "total_params": total,
        "trainable_params": trainable,
        "non_trainable_params": total - trainable,
    }

estimate_model_size_mb

estimate_model_size_mb(model: Module) -> float

Estimates serialized state dictionary size on disk in Megabytes (MB).

Parameters:

Name	Type	Description	Default
`model`	`Module`	PyTorch model.	required

Returns:

Name	Type	Description
`float`	`float`	Estimated file size in MB.

Source code in src/shrinkai/profiler/memory.py

def estimate_model_size_mb(model: nn.Module) -> float:
    """Estimates serialized state dictionary size on disk in Megabytes (MB).

    Args:
        model: PyTorch model.

    Returns:
        float: Estimated file size in MB.
    """
    buffer = io.BytesIO()
    torch.save(model.state_dict(), buffer)
    size_mb = buffer.getbuffer().nbytes / (1024 * 1024)
    return round(size_mb, 2)

get_device_memory_mb

get_device_memory_mb(
    device: device | str = "auto",
) -> float

Returns the currently allocated memory on the specified hardware accelerator.

This is crucial for edge AI profiling, as VRAM is often the primary bottleneck.

Parameters:

Name	Type	Description	Default
`device`	`device \| str`	Computing device ('auto', 'cuda', 'mps', 'cpu', or torch.device).	`'auto'`

Returns:

Name	Type	Description
`float`	`float`	Allocated accelerator memory in MB. Returns 0.0 for CPU (use `get_process_ram_mb` for host CPU RAM instead).

Source code in src/shrinkai/profiler/memory.py

def get_device_memory_mb(device: torch.device | str = "auto") -> float:
    """Returns the currently allocated memory on the specified hardware accelerator.

    This is crucial for edge AI profiling, as VRAM is often the primary bottleneck.

    Args:
        device: Computing device ('auto', 'cuda', 'mps', 'cpu', or torch.device).

    Returns:
        float: Allocated accelerator memory in MB. Returns 0.0 for CPU
            (use `get_process_ram_mb` for host CPU RAM instead).
    """
    resolved_device = resolve_device(device)

    if resolved_device.type == "cuda":
        mem_bytes = torch.cuda.memory_allocated(resolved_device)
        return round(mem_bytes / (1024 * 1024), 2)

    elif resolved_device.type == "mps" and hasattr(torch.mps, "current_allocated_memory"):
        mem_bytes = torch.mps.current_allocated_memory()
        return round(mem_bytes / (1024 * 1024), 2)

    return 0.0

get_process_ram_mb

get_process_ram_mb() -> float

Returns current host RAM consumption of the running Python process.

Returns:

Name	Type	Description
`float`	`float`	Resident memory in MB.

Source code in src/shrinkai/profiler/memory.py

def get_process_ram_mb() -> float:
    """Returns current host RAM consumption of the running Python process.

    Returns:
        float: Resident memory in MB.
    """
    process = psutil.Process(os.getpid())
    return round(process.memory_info().rss / (1024 * 1024), 2)

measure_latency

measure_latency(
    model: Module,
    sample_input: Tensor,
    device: device,
    num_runs: int = 100,
    warmup_runs: int = 15,
) -> dict[str, float]

Measures average inference latency and throughput (FPS / samples per second).

Parameters:

Name	Type	Description	Default
`model`	`Module`	Model to profile.	required
`sample_input`	`Tensor`	Input batch tensor matching inference shape.	required
`device`	`device`	Computing device.	required
`num_runs`	`int`	Number of timed inference iterations.	`100`
`warmup_runs`	`int`	Initial iterations discarded to warm up hardware caches.	`15`

Returns:

Type	Description
`dict[str, float]`	dict[str, float]: Latency per sample (ms), per batch (ms), and throughput (FPS).

Source code in src/shrinkai/profiler/latency.py

def measure_latency(
    model: nn.Module,
    sample_input: torch.Tensor,
    device: torch.device,
    num_runs: int = 100,
    warmup_runs: int = 15,
) -> dict[str, float]:
    """Measures average inference latency and throughput (FPS / samples per second).

    Args:
        model: Model to profile.
        sample_input: Input batch tensor matching inference shape.
        device: Computing device.
        num_runs: Number of timed inference iterations.
        warmup_runs: Initial iterations discarded to warm up hardware caches.

    Returns:
        dict[str, float]: Latency per sample (ms), per batch (ms), and throughput (FPS).
    """
    was_training = model.training
    first_param = next(model.parameters(), None)
    original_device = first_param.device if first_param is not None else None

    model.eval()
    model.to(device)
    sample_input = sample_input.to(device)
    batch_size = sample_input.size(0)

    with torch.no_grad():
        for _ in range(warmup_runs):
            _ = model(sample_input)
        synchronize_device(device)

    timings: list[float] = []
    with torch.no_grad():
        for _ in range(num_runs):
            synchronize_device(device)
            start_time = time.perf_counter()
            _ = model(sample_input)
            synchronize_device(device)
            timings.append(time.perf_counter() - start_time)

    avg_batch_latency_s = sum(timings) / len(timings)
    avg_batch_latency_ms = avg_batch_latency_s * 1000.0
    avg_sample_latency_ms = avg_batch_latency_ms / batch_size
    fps = (batch_size * num_runs) / sum(timings)

    model.train(mode=was_training)
    if original_device is not None:
        model.to(original_device)

    return {
        "batch_latency_ms": avg_batch_latency_ms,
        "sample_latency_ms": avg_sample_latency_ms,
        "fps": fps,
    }

profiler

shrinkai.profiler

Classes

BenchmarkReport

Methods:

show

ModelProfile dataclass

Profiler

Methods:

compare staticmethod

Functions:

count_flops

count_parameters

estimate_model_size_mb

get_device_memory_mb

get_process_ram_mb

measure_latency

ModelProfile `dataclass`

compare `staticmethod`