screenpipe-rfdetr-mlx 源码解析：用纯 Rust + mlx-rs 在 Apple Silicon 上以 126 fps 运行 RF-DETR-Nano 图像 PII 检测器

screenpipe-rfdetr-mlx 源码解析用纯 Rust mlx-rs 在 Apple Silicon 上以 126 fps 运行 RF-DETR-Nano 图像 PII 检测器【免费下载链接】screenpipeYC (S26) | Open Computer History | Record your screen continuously locally and provide context to your agents (Claude, Codex, Openclaw, Hermes, Runner...)项目地址: https://gitcode.com/GitHub_Trending/sc/screenpipescreenpipe 是一款持续录制屏幕并为本地产物Claude、Codex、Openclaw 等 Agent提供上下文的应用而屏幕录制数据中天然包含邮箱、手机号、地址等个人隐私信息PII。为避免这些敏感内容被送入外部模型screenpipe 在本地对画面做 PII 检测与打码。crates/screenpipe-rfdetr-mlx正是这条链路的 Apple Silicon 专用推理引擎它在纯 Rust 中基于mlx-rs复刻了 RF-DETR-Nano 检测器作为screenpipe-redact的mlx-macfeature 的运行时后端。读完本文你将掌握该 crate 的整体架构、模型权重的加载与校验方式、六条可复现的性能优化手法以及移植 ONNX 模型到 MLX 时最容易被忽略的三个陷阱。项目定位图像级 PII 检测的 MLX 运行时screenpipe-rfdetr-mlx是 RF-DETR-Nano 图像 PII 检测器在 Apple SiliconM 系列上的纯 Rust 实现仅面向 macOS aarch64 编译其余平台编译为拒绝运行的桩代码stub。它的直接消费方是screenpipe-redact的mlx-macfeature在 crates/screenpipe-redact/Cargo.toml 中可以看到mlx-mac [dep:screenpipe-rfdetr-mlx]的 feature 声明适配器位于 crates/screenpipe-redact/src/adapters/rfdetr_mlx.rs。模型与权重使用同一份rfdetr_v9检查点同样的模型、同样的权重、同样的图像预处理只是换了一条推理后端。在 8 张图的 PII 语料库M 系列、F32、单图 batch上crate 自带基准对比如下运行时p50 msfps相对 ONNX 加速本 crate纯 MLX7.9126—ONNX Runtime, CPU EP51.3206.5×ONNX Runtime, CoreML EP53.9196.8×同时与 ONNX 的数值一致性验证结果为41/42 个检测框在 IoU ≥ 0.7 下与 ONNX 一一对应最大分数差异仅 0.075。注意这是仓库自带基准的实测值具体数字会随芯片代际M1/M2/M3…、MLX 版本与批次配置变化。目录结构一个最小但完整的 DETR 实现整个 crate 只依赖image图像解码与缩放和thiserror以及 Apple Silicon 目标下条件编译的mlx-rs、mlx-sys、safetensors见 Cargo.toml。源码布局如下src/backbone/ — DINOv2-small ViT12 层 windowed attention隐藏维度 384src/encoder/ — 空目录LWDETR 这里只用了 backbone 的特征金字塔src/decoder/ — 2 层 LWDETR 解码器单尺度 deformable attentionsrc/deformable_attn/ — 纯 MLX 实现的 deformable 交叉注意力src/postprocess.rs — sigmoid 每 query argmax 阈值过滤src/weights.rs — safetensors 权重加载器src/util.rs —mlx_contiguous封装mlx-rs 未暴露该算子通过mlx-sys调用底层 C 接口。入口 src/lib.rs 定义了 13 个类别12 个 PII 类别private_person、private_email、private_phone、private_address、private_url、private_company、private_repo、private_handle、private_channel、private_id、private_date、secret外加no_object类别顺序与训练检查点一致。单条检测结果Detection由像素坐标系下的(x, y, w, h)包围盒、类别索引和置信度分数组成。推理入口load → build → detectRfDetr的公开 API 分三步src/lib.rslet mut model rfdetr_mlx::RfDetr::load(weights/rfdetr_v9.safetensors)?; model.build()?; // 惰性构建 backbone/projector/decoder 子模块 let detections model.detect(image)?; // 默认阈值 0.5其中detect_with_threshold(img, t)允许自定义分数阈值src/lib.rs内部流程为Lanczos3 缩放到 384×384 → 按 ImageNet mean/std[0.485, 0.456, 0.406]/[0.229, 0.224, 0.225]归一化 → 前向得到(boxes, logits)→ 对每 query 的 logits 做 sigmoid 并排除no_object槽位取 argmax → 低于阈值则丢弃 → 将归一化的(cx, cy, w, h)反算回原始图像像素坐标最后按分数降序排列。调试期还可使用forward_profile拿到 backbone / projector / decoder / 总计四段耗时。权重加载safetensors 与架构自校验权重以 safetensors 格式分发加载逻辑在 src/weights.rs先反序列化 header再从元数据中读取model字段并强制要求等于rfdetr_nano_v9不匹配直接拒绝加载——这防止了把其他架构的权重接到本代码上产生静默错误。所有 ONNX initializer 按原始 ONNX 名称如onnx::MatMul_3985存入HashMapString, ArrayRust 侧对命名有唯一话语权各模块forward()按图里出现的名字取权重。模型文件默认路径由适配器约定为~/.screenpipe/models/rfdetr_v9.safetensors适配器首次运行时负责下载并做 SHA-256 校验见 crates/screenpipe-redact/src/adapters/rfdetr_mlx.rs。在 HuggingFace 的screenpipe/pii-image-redactor仓库发布对应的 safetensors 姊妹文件之前你可以用归档仓库screenpipe/rfdetr-mlx中的convert/onnx_to_mlx.py自行从现有rfdetr_v9.onnx转换单次转换、确定性、可复现。FP16/BF16仅环境变量开关默认关闭权重加载时支持通过RFDETR_MLX_FP16环境变量把 F32 权重降精度src/weights.rsRFDETR_MLX_FP16f16或1→ IEEE half在 M1/M2 上表现最好RFDETR_MLX_FP16bf16→ bfloat16在当前硬件上更慢未设置 → 保持 F32默认也是基准数据采用的配置。原因在 README 的性能章节有明确结论Apple GPU 没有原生 BF16而 F16 内核在这些形状下会命中 fallback 路径因此默认 F32低精度只是可选项。开启后forward()会先把输入张量 cast 到权重精度、跑完再 cast 回 F32 交给后处理src/lib.rs。性能笔记六条经过实测的有效优化README 的 “Performance notes — what worked” 记录了移植过程中真正见效的六项优化它们与源码一一对应编码器只做一次尾部 eval。逐层eval()会阻塞 GPU 流水线在 12 个 block 结束后统一 eval 一次让 MLX 自由调度CSE公共子表达式消除会自动去重 4 个多尺度快照共享的依赖。实现见 src/backbone/mod.rs每层 forward 之后只在收集点保存克隆最后统一x.eval()注释明确指出“太频繁的 eval 阻塞 GPU 流水线完全不 eval 会构建 12 层深的递归图导致 MLX 递归 eval 栈溢出”约节省 backbone 端 3ms。把 LayerScale 折叠进 out_proj/fc2 权重。LayerScale 本质是逐输出通道的乘法linear(x, w, b) * lambda linear(x, w*lambda, b*lambda)。在加载期预乘权重与偏置每次前向每个 encoder block 省掉 2 次逐元素乘法。见 src/backbone/encoder_block.rs。常量权重在加载期预转置。解码器自注意力out_proj的权重此前每次前向都要转置并物化现在在加载时转置 contiguous一次src/decoder/mod.rs按注释约省 1.5% 端到端时间。conv 权重只物化一次。PatchEmbed 的 NHWC 权重原本是跨步转置视图conv2d每次调用都会内部拷贝改用crate::util::contiguous在构造期物化src/backbone/dinov2.rs。Array::clone()只是引用计数 1不是深拷贝。所有只读句柄残差捷径、refpoint 张量都用clone()替代原先的deep_clone()见 src/backbone/encoder_block.rs 与 src/decoder/mod.rs。Fused QKV 经历过一次“反转”。把 Q/K/V 拼成单个(384, 1152)matmul 再切分在逐层 eval 时代反而回退约 20%——MLX 当时能把三个小 matmul 并行调度融合反而串行化等 tail-eval 与 LayerScale 折叠落地后每 block 路径变短一次 kernel dispatch 胜过三次融合现在约快 3%。代码里用常量USE_FUSED_QKV true保留了这个开关方便 MLX 再次优化时一行切回src/backbone/encoder_block.rs。移植 ONNX 到 mlx-rs 的三个陷阱README 结尾记录了移植期踩过的三个坑前两个已在源码中形成固定处理范式transpose_axes返回的是跨步视图multiply(one)不会强制物化mlx-rs 的优化器会把*1折叠掉。凡是在 transpose-then-reshape 边界必须调用crate::util::contiguous——它绕过 mlx-rs 未暴露的封装直接经mlx-sys调用原始mlx_contiguousC 算子src/util.rs等价于 PyTorch 的.contiguous()且比add(zeros_of_target_shape)更便宜。ONNX 导出nn.Linear有两种布局3D 输入导出为MatMul权重为(in, out)2D 输入导出为Gemm(transB1)权重为(out, in)。解码器MultiheadAttention.out_proj恰好走 2D 路径权重需要.t()而 in-projection 不需要。所以 backbone 的线性层直接用(in, out)权重src/backbone/encoder_block.rs而解码器out_proj在加载时预转置src/decoder/mod.rs。RF-DETR 解码器的交叉注意力取的是原始的 projector 输出作为 memory而不是enc_output的 LinearLN 结果。enc_output只用于两阶段 proposal head。这一点在 src/decoder/mod.rs 有明确注释与实现decoder layer 循环里传入的是tokens_flat原始扁平化输出而enc_memory仅用于 top-K 选取与 bbox delta 计算。模型内部结构速览从源码可以梳理出端到端的数据流方便后续阅读Embeddingssrc/backbone/dinov2.rsConv2d(3→384, kernel16, stride16)做 patch 嵌入 → 展平 576 个 token → 前置 cls token → 加位置编码(1, 577, 384)→ 按num_windows2把 24×24 网格切成 2×2 共 4 个窗口每个窗口 145 token输出(B*4, 145, 384)。RF-DETR 变体与上游 DINOv2 的两个差异patch16上游 14以及 windowed self-attention。Backbone encodersrc/backbone/mod.rs12 个 pre-norm blockLayerScale 已折叠其中层{3, 6, 9}运行全注意力reshape 回(B, 4*145, 384)再切回其余层保持窗口化在层{2, 5, 8, 11}之后快照隐状态经共享的后编码 LayerNorm输出 4 个尺度的特征金字塔。注释还解释了快照索引与上游配置的 off-by-one 关系。Projector把 4 个尺度的特征拉平投影为(B, 576, 256)的tokens_flat。LWDETR decodersrc/decoder/mod.rsanchor proposals24×24 网格上每位置(cx, cy, 0.05, 0.05)→ 基于enc_memory的两阶段 top-K300 选取 → 学习到的refpoint_embed与 top-K 参考框重参数化 → DAB-DETR 风格 4D 正弦位置编码 → 2 层“自注意力 deformable 交叉注意力 ReLU FFN”的 post-norm 解码层 → 最终 class/bbox head 与 refpoint 精修输出(boxes, logits)对应 ONNX 的(dets, labels)。适用前提与限制本 crate 仅支持 macOS Apple Siliconaarch64其他平台Windows、Linux、Intel Mac会编译为返回Error::UnsupportedTarget的桩实现src/lib.rs保证 workspace 全平台构建不引入 Apple 依赖。性能基准基于 M 系列芯片、F32、单图 batch使用本 crate 自带的 8 图 PII 语料换硬件或换批次配置需重新实测。权重必须先就位~/.screenpipe/models/rfdetr_v9.safetensors且 safetensors 元数据中的model必须为rfdetr_nano_v9。如果你需要在 Apple Silicon 上以接近实时126 fps的速度做本地图像 PII 检测且不希望引入 ONNX Runtime 的运行时负担screenpipe-rfdetr-mlx这份实现连同其性能优化清单是一份可直接复用的工程模板。【免费下载链接】screenpipeYC (S26) | Open Computer History | Record your screen continuously locally and provide context to your agents (Claude, Codex, Openclaw, Hermes, Runner...)项目地址: https://gitcode.com/GitHub_Trending/sc/screenpipe创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考