feat: add optimization toggle flags + benchmark matrix script
use_sdpa, stage_gc flags propagated through backbone/model/pipeline. bench_optimizations.py runs exhaustive/ablation/key sweep of all 6 toggles (slim, stage_gc, sdpa, bf16, fused_decode, gpu_preprocess) with latency + peak memory reporting. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
cmoyates
parent
31d83e218b
commit
d13ba1639d
450
scripts/bench_optimizations.py
Normal file
450
scripts/bench_optimizations.py
Normal file
@ -0,0 +1,450 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Benchmark matrix for wave 2 optimizations.
|
||||
|
||||
Runs an exhaustive matrix of optimization toggles and reports latency + peak
|
||||
memory for each configuration. Helps isolate the impact of each optimization.
|
||||
|
||||
Usage:
|
||||
uv run python scripts/bench_optimizations.py
|
||||
uv run python scripts/bench_optimizations.py --resolution 512 --bench-runs 5
|
||||
uv run python scripts/bench_optimizations.py --checkpoint ckpt.safetensors
|
||||
|
||||
All mx.eval() calls are MLX array materialization, NOT Python eval().
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import itertools
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import mlx.core as mx
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
# Ensure package is importable when running from repo root
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
|
||||
|
||||
from corridorkey_mlx.io.image import preprocess as preprocess_numpy
|
||||
from corridorkey_mlx.io.preprocess_mlx import preprocess_mlx
|
||||
from corridorkey_mlx.model.corridorkey import GreenFormer
|
||||
|
||||
console = Console()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Optimization flags
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
TOGGLE_NAMES = ["slim", "stage_gc", "sdpa", "bf16", "fused_decode", "gpu_preprocess"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class OptConfig:
|
||||
"""Single optimization configuration to benchmark."""
|
||||
|
||||
slim: bool = True
|
||||
stage_gc: bool = True
|
||||
sdpa: bool = True
|
||||
bf16: bool = True
|
||||
fused_decode: bool = True
|
||||
gpu_preprocess: bool = True
|
||||
|
||||
@property
|
||||
def label(self) -> str:
|
||||
flags = []
|
||||
for name in TOGGLE_NAMES:
|
||||
val = getattr(self, name)
|
||||
flags.append(f"{name}={'on' if val else 'off'}")
|
||||
return " | ".join(flags)
|
||||
|
||||
@property
|
||||
def short_label(self) -> str:
|
||||
parts = []
|
||||
for name in TOGGLE_NAMES:
|
||||
if getattr(self, name):
|
||||
parts.append(name)
|
||||
return "+".join(parts) if parts else "baseline"
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchResult:
|
||||
"""Result from a single benchmark configuration."""
|
||||
|
||||
config: OptConfig
|
||||
resolution: int
|
||||
warmup_ms: float = 0.0
|
||||
median_ms: float = 0.0
|
||||
min_ms: float = 0.0
|
||||
peak_memory_mb: float = 0.0
|
||||
all_times: list[float] = field(default_factory=list)
|
||||
error: str | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark logic
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_model(
|
||||
config: OptConfig,
|
||||
img_size: int,
|
||||
checkpoint: Path | None,
|
||||
) -> GreenFormer:
|
||||
"""Build model with the given optimization config."""
|
||||
dtype = mx.bfloat16 if config.bf16 else mx.float32
|
||||
model = GreenFormer(
|
||||
img_size=img_size,
|
||||
dtype=dtype,
|
||||
fused_decode=config.fused_decode,
|
||||
slim=config.slim,
|
||||
use_sdpa=config.sdpa,
|
||||
stage_gc=config.stage_gc,
|
||||
)
|
||||
if checkpoint and checkpoint.exists():
|
||||
model.load_checkpoint(checkpoint)
|
||||
else:
|
||||
model.eval()
|
||||
mx.eval(model.parameters()) # noqa: S307
|
||||
return model
|
||||
|
||||
|
||||
def make_input(img_size: int, gpu_preprocess: bool) -> mx.array:
|
||||
"""Create input tensor, optionally via GPU preprocessing path."""
|
||||
import numpy as np
|
||||
|
||||
np.random.seed(42)
|
||||
rgb_f32 = np.random.rand(img_size, img_size, 3).astype(np.float32)
|
||||
mask_f32 = np.random.rand(img_size, img_size, 1).astype(np.float32)
|
||||
|
||||
if gpu_preprocess:
|
||||
return preprocess_mlx(rgb_f32, mask_f32)
|
||||
return preprocess_numpy(rgb_f32, mask_f32)
|
||||
|
||||
|
||||
def time_inference(
|
||||
model: GreenFormer,
|
||||
x: mx.array,
|
||||
warmup_runs: int,
|
||||
bench_runs: int,
|
||||
) -> tuple[float, float, float, list[float]]:
|
||||
"""Run warmup + benchmark, return (warmup_ms, median_ms, min_ms, all_times)."""
|
||||
|
||||
def run() -> None:
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
# Warmup
|
||||
start = time.perf_counter()
|
||||
run()
|
||||
warmup_ms = (time.perf_counter() - start) * 1000.0
|
||||
|
||||
for _ in range(warmup_runs - 1):
|
||||
run()
|
||||
|
||||
# Benchmark
|
||||
times: list[float] = []
|
||||
for _ in range(bench_runs):
|
||||
start = time.perf_counter()
|
||||
run()
|
||||
times.append((time.perf_counter() - start) * 1000.0)
|
||||
|
||||
times.sort()
|
||||
median_ms = times[len(times) // 2]
|
||||
min_ms = times[0]
|
||||
return warmup_ms, median_ms, min_ms, times
|
||||
|
||||
|
||||
def measure_peak_memory(model: GreenFormer, x: mx.array) -> float:
|
||||
"""Measure peak Metal memory for a single forward pass (MB)."""
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
# Use non-deprecated API if available, fall back to mx.metal.*
|
||||
reset_fn = getattr(mx, "reset_peak_memory", None) or getattr(
|
||||
mx.metal, "reset_peak_memory", None
|
||||
)
|
||||
get_fn = getattr(mx, "get_peak_memory", None) or getattr(mx.metal, "get_peak_memory", None)
|
||||
if not reset_fn or not get_fn:
|
||||
return 0.0
|
||||
|
||||
reset_fn()
|
||||
|
||||
out = model(x)
|
||||
mx.eval(out) # noqa: S307
|
||||
|
||||
peak_bytes = get_fn()
|
||||
|
||||
del out
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
return peak_bytes / (1024 * 1024)
|
||||
|
||||
|
||||
def bench_config(
|
||||
config: OptConfig,
|
||||
resolution: int,
|
||||
checkpoint: Path | None,
|
||||
warmup_runs: int,
|
||||
bench_runs: int,
|
||||
) -> BenchResult:
|
||||
"""Benchmark a single optimization configuration."""
|
||||
result = BenchResult(config=config, resolution=resolution)
|
||||
|
||||
try:
|
||||
model = build_model(config, resolution, checkpoint)
|
||||
x = make_input(resolution, config.gpu_preprocess)
|
||||
mx.eval(x) # noqa: S307
|
||||
|
||||
warmup_ms, median_ms, min_ms, all_times = time_inference(model, x, warmup_runs, bench_runs)
|
||||
result.warmup_ms = warmup_ms
|
||||
result.median_ms = median_ms
|
||||
result.min_ms = min_ms
|
||||
result.all_times = all_times
|
||||
|
||||
# Measure peak memory with a fresh model to avoid cached state
|
||||
model_fresh = build_model(config, resolution, checkpoint)
|
||||
result.peak_memory_mb = measure_peak_memory(model_fresh, x)
|
||||
|
||||
# Cleanup
|
||||
del model, model_fresh, x
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Matrix generation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def generate_matrix(sweep: str) -> list[OptConfig]:
|
||||
"""Generate benchmark configurations.
|
||||
|
||||
Args:
|
||||
sweep: "full" for exhaustive 2^6 matrix, "ablation" for toggle-one-off,
|
||||
"key" for important combinations only.
|
||||
"""
|
||||
all_on = OptConfig()
|
||||
|
||||
if sweep == "ablation":
|
||||
# All-on baseline + toggle each flag off one at a time
|
||||
configs = [all_on]
|
||||
for name in TOGGLE_NAMES:
|
||||
off = OptConfig(**{**vars(all_on), name: False})
|
||||
configs.append(off)
|
||||
# Also add all-off baseline
|
||||
configs.append(
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
)
|
||||
)
|
||||
return configs
|
||||
|
||||
if sweep == "full":
|
||||
# Exhaustive 2^6 = 64 configurations
|
||||
configs = []
|
||||
for combo in itertools.product([True, False], repeat=len(TOGGLE_NAMES)):
|
||||
kwargs = dict(zip(TOGGLE_NAMES, combo, strict=True))
|
||||
configs.append(OptConfig(**kwargs))
|
||||
return configs
|
||||
|
||||
# "key" — important configurations
|
||||
return [
|
||||
# All off (wave 1 baseline)
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Wave 1 only (bf16 + fused_decode)
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=False,
|
||||
bf16=True,
|
||||
fused_decode=True,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Wave 2 only (slim + stage_gc + sdpa + gpu_preprocess)
|
||||
OptConfig(
|
||||
slim=True,
|
||||
stage_gc=True,
|
||||
sdpa=True,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=True,
|
||||
),
|
||||
# All on
|
||||
all_on,
|
||||
# SDPA only
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=False,
|
||||
sdpa=True,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
# Stage GC only
|
||||
OptConfig(
|
||||
slim=False,
|
||||
stage_gc=True,
|
||||
sdpa=False,
|
||||
bf16=False,
|
||||
fused_decode=False,
|
||||
gpu_preprocess=False,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Output
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_results(results: list[BenchResult]) -> None:
|
||||
"""Print results as a rich table."""
|
||||
table = Table(title="Optimization Benchmark Matrix")
|
||||
|
||||
table.add_column("Config", max_width=50)
|
||||
table.add_column("Res", justify="right")
|
||||
table.add_column("Warmup (ms)", justify="right")
|
||||
table.add_column("Median (ms)", justify="right")
|
||||
table.add_column("Min (ms)", justify="right")
|
||||
table.add_column("Peak Mem (MB)", justify="right")
|
||||
table.add_column("Status", justify="center")
|
||||
|
||||
# Sort by resolution then median time
|
||||
results.sort(key=lambda r: (r.resolution, r.median_ms))
|
||||
|
||||
for r in results:
|
||||
if r.error:
|
||||
table.add_row(
|
||||
r.config.short_label,
|
||||
str(r.resolution),
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
f"[red]FAIL: {r.error[:30]}[/red]",
|
||||
)
|
||||
else:
|
||||
table.add_row(
|
||||
r.config.short_label,
|
||||
str(r.resolution),
|
||||
f"{r.warmup_ms:.1f}",
|
||||
f"{r.median_ms:.1f}",
|
||||
f"{r.min_ms:.1f}",
|
||||
f"{r.peak_memory_mb:.0f}" if r.peak_memory_mb > 0 else "N/A",
|
||||
"[green]OK[/green]",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
# Print speedup summary vs baseline
|
||||
baselines = {r.resolution: r for r in results if r.config.short_label == "baseline"}
|
||||
if baselines:
|
||||
console.print("\n[bold]Speedup vs baseline (all-off):[/bold]")
|
||||
for r in results:
|
||||
if r.error or r.config.short_label == "baseline":
|
||||
continue
|
||||
baseline = baselines.get(r.resolution)
|
||||
if baseline and baseline.median_ms > 0:
|
||||
speedup = baseline.median_ms / r.median_ms
|
||||
mem_ratio = (
|
||||
f"{r.peak_memory_mb / baseline.peak_memory_mb:.2f}x"
|
||||
if baseline.peak_memory_mb > 0 and r.peak_memory_mb > 0
|
||||
else "N/A"
|
||||
)
|
||||
console.print(
|
||||
f" {r.config.short_label:40s} "
|
||||
f"{speedup:.2f}x speed, {mem_ratio} memory "
|
||||
f"({r.resolution})"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Benchmark matrix for wave 2 optimizations")
|
||||
parser.add_argument(
|
||||
"--checkpoint",
|
||||
type=Path,
|
||||
default=Path("checkpoints/corridorkey_mlx.safetensors"),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[512],
|
||||
help="Resolution(s) to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sweep",
|
||||
choices=["key", "ablation", "full"],
|
||||
default="ablation",
|
||||
help="key=6 configs, ablation=8 configs (toggle one off), full=64 configs",
|
||||
)
|
||||
parser.add_argument("--warmup-runs", type=int, default=3)
|
||||
parser.add_argument("--bench-runs", type=int, default=10)
|
||||
args = parser.parse_args()
|
||||
|
||||
ckpt = args.checkpoint if args.checkpoint.exists() else None
|
||||
configs = generate_matrix(args.sweep)
|
||||
|
||||
console.print("[bold]CorridorKey MLX — Optimization Benchmark Matrix[/bold]")
|
||||
console.print(f" Weights: {args.checkpoint if ckpt else 'random (no checkpoint)'}")
|
||||
console.print(f" Resolutions: {args.resolution}")
|
||||
console.print(f" Sweep: {args.sweep} ({len(configs)} configs)")
|
||||
console.print(f" Warmup: {args.warmup_runs}, Bench: {args.bench_runs}")
|
||||
console.print()
|
||||
|
||||
all_results: list[BenchResult] = []
|
||||
total = len(configs) * len(args.resolution)
|
||||
idx = 0
|
||||
|
||||
for res in args.resolution:
|
||||
for config in configs:
|
||||
idx += 1
|
||||
console.print(
|
||||
f" [{idx}/{total}] {res}x{res} — {config.short_label}...",
|
||||
end=" ",
|
||||
)
|
||||
result = bench_config(config, res, ckpt, args.warmup_runs, args.bench_runs)
|
||||
if result.error:
|
||||
console.print("[red]FAIL[/red]")
|
||||
else:
|
||||
console.print(
|
||||
f"[green]{result.median_ms:.1f}ms[/green] "
|
||||
f"(peak: {result.peak_memory_mb:.0f}MB)"
|
||||
if result.peak_memory_mb > 0
|
||||
else f"[green]{result.median_ms:.1f}ms[/green]"
|
||||
)
|
||||
all_results.append(result)
|
||||
|
||||
console.print()
|
||||
print_results(all_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -36,6 +36,8 @@ def load_model(
|
||||
dtype: mx.Dtype = mx.bfloat16,
|
||||
fused_decode: bool = True,
|
||||
slim: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
stage_gc: bool = True,
|
||||
) -> GreenFormer:
|
||||
"""Build GreenFormer and load weights from safetensors checkpoint.
|
||||
|
||||
@ -53,8 +55,17 @@ def load_model(
|
||||
Metal dispatch calls. Bit-exact with unfused path.
|
||||
slim: If True, forward returns only 4 final keys (drops intermediates).
|
||||
Reduces reference lifetime so MLX can reclaim buffers sooner.
|
||||
use_sdpa: If True, use mx.fast.scaled_dot_product_attention in backbone.
|
||||
stage_gc: If True, materialize + GC at backbone/decoder/refiner boundaries.
|
||||
"""
|
||||
model = GreenFormer(img_size=img_size, dtype=dtype, fused_decode=fused_decode, slim=slim)
|
||||
model = GreenFormer(
|
||||
img_size=img_size,
|
||||
dtype=dtype,
|
||||
fused_decode=fused_decode,
|
||||
slim=slim,
|
||||
use_sdpa=use_sdpa,
|
||||
stage_gc=stage_gc,
|
||||
)
|
||||
model.load_checkpoint(checkpoint)
|
||||
if compile:
|
||||
model = compile_model(model, shapeless=shapeless)
|
||||
|
||||
@ -38,13 +38,16 @@ class GreenFormer(nn.Module):
|
||||
dtype: mx.Dtype = mx.float32,
|
||||
fused_decode: bool = False,
|
||||
slim: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
stage_gc: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._compute_dtype = dtype
|
||||
self._fused_decode = fused_decode
|
||||
self._slim = slim
|
||||
self._stage_gc = stage_gc
|
||||
self._compiled = False
|
||||
self.backbone = HieraBackbone(img_size=img_size)
|
||||
self.backbone = HieraBackbone(img_size=img_size, use_sdpa=use_sdpa)
|
||||
self.alpha_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=1)
|
||||
self.fg_decoder = DecoderHead(BACKBONE_CHANNELS, EMBED_DIM, output_dim=3)
|
||||
self.refiner = CNNRefinerModule()
|
||||
@ -73,7 +76,7 @@ class GreenFormer(nn.Module):
|
||||
|
||||
# Materialize backbone output so MLX can free intermediate graph nodes.
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
if not self._compiled:
|
||||
if self._stage_gc and not self._compiled:
|
||||
mx.eval(features) # noqa: S307
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
@ -104,7 +107,7 @@ class GreenFormer(nn.Module):
|
||||
|
||||
# Materialize decoder output so MLX can free decoder graph nodes.
|
||||
# NOTE: mx.eval is MLX array materialization, not Python eval()
|
||||
if not self._compiled:
|
||||
if self._stage_gc and not self._compiled:
|
||||
mx.eval(alpha_coarse, fg_coarse, alpha_logits_up, fg_logits_up) # noqa: S307
|
||||
gc.collect()
|
||||
mx.clear_cache()
|
||||
|
||||
@ -251,6 +251,7 @@ class MaskUnitAttention(nn.Module):
|
||||
q_stride: int = 1,
|
||||
window_size: int = 0,
|
||||
use_mask_unit_attn: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dim_out = dim_out
|
||||
@ -260,6 +261,7 @@ class MaskUnitAttention(nn.Module):
|
||||
self.scale = self.head_dim**-0.5
|
||||
self.window_size = window_size
|
||||
self.use_mask_unit_attn = use_mask_unit_attn
|
||||
self._use_sdpa = use_sdpa
|
||||
|
||||
self.qkv = nn.Linear(dim, 3 * dim_out)
|
||||
self.proj = nn.Linear(dim_out, dim_out)
|
||||
@ -284,6 +286,7 @@ class MaskUnitAttention(nn.Module):
|
||||
q = q.reshape(batch_size, self.heads, num_windows, self.q_stride, -1, self.head_dim)
|
||||
q = mx.max(q, axis=3)
|
||||
|
||||
if self._use_sdpa:
|
||||
# Fold window dim into batch for SDPA (expects 4D).
|
||||
# Q/K/V are (B, heads, windows, tokens, head_dim) — transpose windows
|
||||
# next to batch before reshape so (B, W, H, T, D) -> (B*W, H, T, D).
|
||||
@ -304,6 +307,11 @@ class MaskUnitAttention(nn.Module):
|
||||
# Unfold: (B*W, H, T, D) -> (B, W, H, T, D) -> (B, H, W, T, D)
|
||||
x = x.reshape(batch_size, num_windows, self.heads, q_tokens, self.head_dim)
|
||||
x = mx.transpose(x, axes=(0, 2, 1, 3, 4))
|
||||
else:
|
||||
# Manual scaled dot-product attention (original implementation)
|
||||
attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3))
|
||||
attn = mx.softmax(attn, axis=-1)
|
||||
x = attn @ v
|
||||
|
||||
# [B, heads, num_windows, tokens, head_dim] -> [B, tokens, num_windows, heads, head_dim]
|
||||
# -> [B, N', dim_out] (matches PyTorch transpose(1, 3))
|
||||
@ -328,6 +336,7 @@ class HieraBlock(nn.Module):
|
||||
q_stride: int = 1,
|
||||
window_size: int = 0,
|
||||
use_mask_unit_attn: bool = False,
|
||||
use_sdpa: bool = True,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
@ -340,7 +349,7 @@ class HieraBlock(nn.Module):
|
||||
self.proj = nn.Linear(dim, dim_out)
|
||||
|
||||
self.attn = MaskUnitAttention(
|
||||
dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn
|
||||
dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn, use_sdpa=use_sdpa
|
||||
)
|
||||
|
||||
self.norm2 = nn.LayerNorm(dim_out)
|
||||
@ -369,7 +378,7 @@ class HieraBackbone(nn.Module):
|
||||
Hardcoded for hiera_base_plus_224 config with 4-channel input.
|
||||
"""
|
||||
|
||||
def __init__(self, img_size: int = 512) -> None:
|
||||
def __init__(self, img_size: int = 512, use_sdpa: bool = True) -> None:
|
||||
super().__init__()
|
||||
self.img_size = img_size
|
||||
|
||||
@ -434,6 +443,7 @@ class HieraBackbone(nn.Module):
|
||||
q_stride=(flat_q_stride if i in q_pool_blocks else 1),
|
||||
window_size=flat_mu_size,
|
||||
use_mask_unit_attn=use_mu_attn,
|
||||
use_sdpa=use_sdpa,
|
||||
)
|
||||
self.blocks.append(block)
|
||||
embed_dim = dim_out
|
||||
|
||||
Loading…
Reference in New Issue
Block a user