lhk229
9a9084b3d6
Three optimizations from profiling the cross-check-dominated pipeline (9700X CPU, all pilot-validated on TestImage3/FixImage1): - Reuse the cross-check HR-matting@2048 forward as the segmentation mask (cross_check.reuse_as_seg, default ON; --no-cross-check-as-seg to opt out). Skips the BiRefNet@1024 load+forward entirely: ~66s -> ~45s, one less 0.9GB model. Trimap 99.8% identical, no structural change. - --precision bf16 now fans out to all three models: ViTMatte keeps its weight cast; both BiRefNets run their forward under autocast with a dispatcher-level AutocastCPU fp32 shim for torchvision::deform_conv2d (no bf16 CPU kernel, no autocast wrapper upstream). Shared hardware gate in bgfilter/precision.py falls back to fp32 off native-bf16 hardware. TestImage3: 51.9s -> 33.1s; alpha diff max 0.15, none >0.25. - MIMALLOC_PURGE_DELAY=0 (bgfilter/__init__.py, before torch loads): Windows torch's bundled mimalloc lazily retains ~10GB of freed BiRefNet activations, stacking under ViTMatte's attention peak. 2048x2048 bf16: peak 25.2 -> 21.4GB and slightly faster (73 -> 63s). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
74 lines
3.1 KiB
Python
74 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import warnings
|
|
|
|
|
|
def resolve_dtype(torch, device, precision: str):
|
|
"""Map a precision string to a torch dtype, gated on hardware support.
|
|
|
|
bf16 is only worth it on hardware with native kernels (any modern GPU, or a
|
|
CPU with AVX512-BF16/AMX). Without them PyTorch silently routes bf16 matmuls
|
|
onto a fallback that is orders of magnitude slower than fp32, so this gate
|
|
falls back to fp32 with a warning instead. The CPU check is the same one
|
|
PyTorch uses to route bf16 matmuls onto oneDNN's native kernels.
|
|
"""
|
|
if precision == "fp32":
|
|
return torch.float32
|
|
if precision != "bf16":
|
|
raise RuntimeError(f"Unsupported precision '{precision}'. Use 'fp32' or 'bf16'.")
|
|
if device.type == "cuda":
|
|
if torch.cuda.is_bf16_supported():
|
|
return torch.bfloat16
|
|
reason = "this GPU has no bf16 support"
|
|
else:
|
|
try:
|
|
native = torch.ops.mkldnn._is_mkldnn_bf16_supported()
|
|
except (AttributeError, RuntimeError):
|
|
native = False
|
|
if native:
|
|
return torch.bfloat16
|
|
reason = "this CPU has no native bf16 support (needs AVX512-BF16/AMX)"
|
|
warnings.warn(f"precision 'bf16' requested but {reason}; running fp32 instead.")
|
|
return torch.float32
|
|
|
|
|
|
# Keep the Library object alive: dropping it would un-register the impl.
|
|
_deform_conv2d_shim: object | None = None
|
|
|
|
|
|
def ensure_deform_conv2d_autocast_cpu(torch) -> None:
|
|
"""Give torchvision's deform_conv2d an AutocastCPU wrapper (fp32 fallback).
|
|
|
|
torchvision (<= 0.27 at least) ships no AutocastCPU registration and no
|
|
BFloat16 CPU kernel for ``torchvision::deform_conv2d``, so running a model
|
|
that uses it (BiRefNet) under ``torch.autocast("cpu", bf16)`` crashes with
|
|
"deformable_im2col not implemented for 'BFloat16'". This registers a
|
|
dispatcher-level wrapper that casts the op's inputs to fp32 and runs it
|
|
outside autocast — upstream/downstream convs keep their bf16 speed, only
|
|
this one op pays a cast. No-op if called twice or if a future torchvision
|
|
registers its own wrapper.
|
|
"""
|
|
global _deform_conv2d_shim
|
|
if _deform_conv2d_shim is not None:
|
|
return
|
|
import torchvision # noqa: F401 (registers torchvision::deform_conv2d)
|
|
|
|
lib = torch.library.Library("torchvision", "IMPL", "AutocastCPU")
|
|
|
|
def _deform_conv2d_fp32(input, weight, offset, mask, bias,
|
|
stride_h, stride_w, pad_h, pad_w,
|
|
dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask):
|
|
with torch.autocast("cpu", enabled=False):
|
|
return torch.ops.torchvision.deform_conv2d(
|
|
input.float(), weight.float(), offset.float(), mask.float(),
|
|
bias.float(), stride_h, stride_w, pad_h, pad_w,
|
|
dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask,
|
|
)
|
|
|
|
try:
|
|
lib.impl("deform_conv2d", _deform_conv2d_fp32)
|
|
except RuntimeError:
|
|
# torchvision grew its own AutocastCPU wrapper; ours is unneeded.
|
|
return
|
|
_deform_conv2d_shim = lib
|