Skip to content

memory_gauge

sleap_nn.config_generator.tui.widgets.memory_gauge

Memory estimation gauge widget.

Displays GPU and CPU memory estimates with color-coded status indicators.

Classes:

Name Description
MemoryBreakdownCard

Detailed memory breakdown card with component-level estimates.
MemoryGauge

Widget displaying memory estimation with color-coded status.

MemoryBreakdownCard

Bases: Static

Detailed memory breakdown card with component-level estimates.

Shows individual memory consumption for: - Model parameters/weights - Batch images - Feature maps/activations - Confidence maps/PAFs - Gradients

Methods:

Name Description
__init__

Initialize the breakdown card.
render

Render the breakdown card.
update_estimate

Update the displayed estimate.
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
115
116
117
118
119
120
121
122
123
124
125
126
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
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
class MemoryBreakdownCard(Static):
    """Detailed memory breakdown card with component-level estimates.

    Shows individual memory consumption for:
    - Model parameters/weights
    - Batch images
    - Feature maps/activations
    - Confidence maps/PAFs
    - Gradients
    """

    DEFAULT_CSS = """
    MemoryBreakdownCard {
        height: auto;
        padding: 1;
        border: solid $surface-lighten-2;
        margin: 1 0;
    }

    MemoryBreakdownCard .card-title {
        text-style: bold;
        margin-bottom: 1;
    }

    MemoryBreakdownCard .memory-row {
        height: 1;
    }

    MemoryBreakdownCard .memory-label {
        width: 1fr;
        color: $text-muted;
    }

    MemoryBreakdownCard .memory-value {
        width: auto;
        text-align: right;
    }
    """

    def __init__(
        self,
        estimate: Optional[MemoryEstimate] = None,
        model_type: str = "Single Instance",
        **kwargs,
    ):
        """Initialize the breakdown card.

        Args:
            estimate: Memory estimate to display.
            model_type: Name of model type for header display.
            **kwargs: Additional arguments passed to parent.
        """
        super().__init__(**kwargs)
        self._estimate = estimate
        self._model_type = model_type

    def update_estimate(
        self, estimate: MemoryEstimate, model_type: Optional[str] = None
    ) -> None:
        """Update the displayed estimate.

        Args:
            estimate: New memory estimate.
            model_type: Optional new model type name.
        """
        self._estimate = estimate
        if model_type:
            self._model_type = model_type
        self.refresh()

    def render(self) -> str:
        """Render the breakdown card."""
        if self._estimate is None:
            return "Memory Breakdown\n" "Calculating..."

        est = self._estimate
        lines = [
            f"GPU Memory ({self._model_type})",
            "─" * 30,
        ]

        # Format values
        params_mb = est.model_memory_gb * 1024
        batch_mb = (est.batch_size * est.image_bytes) / (1024 * 1024)
        activations_mb = est.activations_memory_gb * 1024
        gradients_mb = est.gradients_memory_gb * 1024

        rows = [
            ("Model Params", f"~{self._format_count(params_mb / 4 * 1e6)}"),
            ("Weights", f"~{params_mb:.0f} MB"),
            ("Batch Images", f"~{batch_mb:.0f} MB"),
            ("Feature Maps", f"~{activations_mb:.0f} MB"),
            ("Gradients", f"~{gradients_mb:.0f} MB"),
        ]

        for label, value in rows:
            lines.append(f"  {label:<16} {value:>10}")

        lines.append("")
        lines.append(f"  {'Total':<16} {est.total_gpu_gb:.1f} GB")

        # Status line
        status_tags = {
            "green": ("[green]", "[/green]"),
            "yellow": ("[yellow]", "[/yellow]"),
            "red": ("[red]", "[/red]"),
        }
        open_tag, close_tag = status_tags.get(est.gpu_status, ("", ""))
        lines.append("")
        lines.append(f"{open_tag}{est.gpu_message}{close_tag}")

        return "\n".join(lines)

    def _format_count(self, count: float) -> str:
        """Format parameter count with appropriate suffix."""
        if count >= 1e9:
            return f"{count/1e9:.1f}B"
        elif count >= 1e6:
            return f"{count/1e6:.1f}M"
        elif count >= 1e3:
            return f"{count/1e3:.1f}K"
        return f"{count:.0f}"

__init__(estimate=None, model_type='Single Instance', **kwargs)

Initialize the breakdown card.

Parameters:

Name Type Description Default
estimate Optional[MemoryEstimate]

Memory estimate to display.

 None
model_type str

Name of model type for header display.

 'Single Instance'
**kwargs

Additional arguments passed to parent.

 {}
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def __init__(
    self,
    estimate: Optional[MemoryEstimate] = None,
    model_type: str = "Single Instance",
    **kwargs,
):
    """Initialize the breakdown card.

    Args:
        estimate: Memory estimate to display.
        model_type: Name of model type for header display.
        **kwargs: Additional arguments passed to parent.
    """
    super().__init__(**kwargs)
    self._estimate = estimate
    self._model_type = model_type

render()

Render the breakdown card.

Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
def render(self) -> str:
    """Render the breakdown card."""
    if self._estimate is None:
        return "Memory Breakdown\n" "Calculating..."

    est = self._estimate
    lines = [
        f"GPU Memory ({self._model_type})",
        "─" * 30,
    ]

    # Format values
    params_mb = est.model_memory_gb * 1024
    batch_mb = (est.batch_size * est.image_bytes) / (1024 * 1024)
    activations_mb = est.activations_memory_gb * 1024
    gradients_mb = est.gradients_memory_gb * 1024

    rows = [
        ("Model Params", f"~{self._format_count(params_mb / 4 * 1e6)}"),
        ("Weights", f"~{params_mb:.0f} MB"),
        ("Batch Images", f"~{batch_mb:.0f} MB"),
        ("Feature Maps", f"~{activations_mb:.0f} MB"),
        ("Gradients", f"~{gradients_mb:.0f} MB"),
    ]

    for label, value in rows:
        lines.append(f"  {label:<16} {value:>10}")

    lines.append("")
    lines.append(f"  {'Total':<16} {est.total_gpu_gb:.1f} GB")

    # Status line
    status_tags = {
        "green": ("[green]", "[/green]"),
        "yellow": ("[yellow]", "[/yellow]"),
        "red": ("[red]", "[/red]"),
    }
    open_tag, close_tag = status_tags.get(est.gpu_status, ("", ""))
    lines.append("")
    lines.append(f"{open_tag}{est.gpu_message}{close_tag}")

    return "\n".join(lines)

update_estimate(estimate, model_type=None)

Update the displayed estimate.

Parameters:

Name Type Description Default
estimate MemoryEstimate

New memory estimate.

 required
model_type Optional[str]

Optional new model type name.

 None
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
def update_estimate(
    self, estimate: MemoryEstimate, model_type: Optional[str] = None
) -> None:
    """Update the displayed estimate.

    Args:
        estimate: New memory estimate.
        model_type: Optional new model type name.
    """
    self._estimate = estimate
    if model_type:
        self._model_type = model_type
    self.refresh()

MemoryGauge

Bases: Static

Widget displaying memory estimation with color-coded status.

Shows estimated GPU and CPU memory usage with visual indicators for whether the configuration will fit on available hardware.

Methods:

Name Description
__init__

Initialize with optional memory estimate.
render

Render the memory gauge display.
toggle_breakdown

Toggle the detailed breakdown visibility.
update_estimate

Update the displayed memory estimate.

Attributes:

Name Type Description
estimate Optional[MemoryEstimate]

Get the current memory estimate.
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 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
class MemoryGauge(Static):
    """Widget displaying memory estimation with color-coded status.

    Shows estimated GPU and CPU memory usage with visual indicators
    for whether the configuration will fit on available hardware.
    """

    DEFAULT_CSS = """
    MemoryGauge {
        height: auto;
        padding: 1;
        border: solid $primary;
        margin-bottom: 1;
    }

    MemoryGauge .memory-title {
        text-style: bold;
    }

    MemoryGauge .status-green {
        color: $success;
    }

    MemoryGauge .status-yellow {
        color: $warning;
    }

    MemoryGauge .status-red {
        color: $error;
    }

    MemoryGauge .memory-breakdown {
        color: $text-muted;
        margin-top: 1;
    }
    """

    def __init__(
        self,
        estimate: Optional[MemoryEstimate] = None,
        show_breakdown: bool = True,
        **kwargs,
    ):
        """Initialize with optional memory estimate.

        Args:
            estimate: Initial memory estimate to display.
            show_breakdown: Whether to show detailed breakdown.
            **kwargs: Additional arguments passed to parent.
        """
        super().__init__(**kwargs)
        self._estimate = estimate
        self._show_breakdown = show_breakdown

    def update_estimate(self, estimate: MemoryEstimate) -> None:
        """Update the displayed memory estimate.

        Args:
            estimate: New memory estimate to display.
        """
        self._estimate = estimate
        self.refresh()

    def toggle_breakdown(self) -> None:
        """Toggle the detailed breakdown visibility."""
        self._show_breakdown = not self._show_breakdown
        self.refresh()

    @property
    def estimate(self) -> Optional[MemoryEstimate]:
        """Get the current memory estimate."""
        return self._estimate

    def render(self) -> str:
        """Render the memory gauge display."""
        if self._estimate is None:
            return "Memory Estimate\n" "──────────────\n" "Loading..."

        status_icons = {"green": "✓", "yellow": "⚠", "red": "✗"}
        icon = status_icons.get(self._estimate.gpu_status, "?")

        lines = [
            "Memory Estimate",
            "──────────────",
            f"GPU: {self._estimate.total_gpu_gb:.1f} GB {icon}",
            f"  {self._estimate.gpu_message}",
            "",
            f"CPU: {self._estimate.cache_memory_gb:.1f} GB",
            f"  {self._estimate.cpu_message}",
        ]

        if self._show_breakdown and hasattr(self._estimate, "breakdown"):
            lines.append("")
            lines.append("Breakdown:")
            breakdown = self._estimate.breakdown
            if breakdown:
                for key, value in breakdown.items():
                    lines.append(f"  {key}: {value:.1f} MB")

        return "\n".join(lines)

estimate property

Get the current memory estimate.

__init__(estimate=None, show_breakdown=True, **kwargs)

Initialize with optional memory estimate.

Parameters:

Name Type Description Default
estimate Optional[MemoryEstimate]

Initial memory estimate to display.

 None
show_breakdown bool

Whether to show detailed breakdown.

 True
**kwargs

Additional arguments passed to parent.

 {}
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def __init__(
    self,
    estimate: Optional[MemoryEstimate] = None,
    show_breakdown: bool = True,
    **kwargs,
):
    """Initialize with optional memory estimate.

    Args:
        estimate: Initial memory estimate to display.
        show_breakdown: Whether to show detailed breakdown.
        **kwargs: Additional arguments passed to parent.
    """
    super().__init__(**kwargs)
    self._estimate = estimate
    self._show_breakdown = show_breakdown

render()

Render the memory gauge display.

Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
 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
def render(self) -> str:
    """Render the memory gauge display."""
    if self._estimate is None:
        return "Memory Estimate\n" "──────────────\n" "Loading..."

    status_icons = {"green": "✓", "yellow": "⚠", "red": "✗"}
    icon = status_icons.get(self._estimate.gpu_status, "?")

    lines = [
        "Memory Estimate",
        "──────────────",
        f"GPU: {self._estimate.total_gpu_gb:.1f} GB {icon}",
        f"  {self._estimate.gpu_message}",
        "",
        f"CPU: {self._estimate.cache_memory_gb:.1f} GB",
        f"  {self._estimate.cpu_message}",
    ]

    if self._show_breakdown and hasattr(self._estimate, "breakdown"):
        lines.append("")
        lines.append("Breakdown:")
        breakdown = self._estimate.breakdown
        if breakdown:
            for key, value in breakdown.items():
                lines.append(f"  {key}: {value:.1f} MB")

    return "\n".join(lines)

toggle_breakdown()

Toggle the detailed breakdown visibility.

Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
76
77
78
79
def toggle_breakdown(self) -> None:
    """Toggle the detailed breakdown visibility."""
    self._show_breakdown = not self._show_breakdown
    self.refresh()

update_estimate(estimate)

Update the displayed memory estimate.

Parameters:

Name Type Description Default
estimate MemoryEstimate

New memory estimate to display.

 required
Source code in sleap_nn/config_generator/tui/widgets/memory_gauge.py 
67
68
69
70
71
72
73
74
def update_estimate(self, estimate: MemoryEstimate) -> None:
    """Update the displayed memory estimate.

    Args:
        estimate: New memory estimate to display.
    """
    self._estimate = estimate
    self.refresh()