Skip to content

benchmark

shrinkai.profiler.benchmark

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.

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 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
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 
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
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 
14
15
16
17
18
19
20
21
22
23
24
@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 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
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 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
@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: