VisionXLab
/

Point2RBox-v3-jittor

@@ -1,173 +1,125 @@
-# torch → Jittor 移植实战笔记（PLAN §8 的实测版）
-> 全部条目都在本仓库实际踩过/验证过，L1/L2 测试锁定。给后来者（含 Agent B）参考。
-## 🔴 会静默出错的（最危险，症状不指向根因）
-| # | 问题 | 症状 | 对策 |
-|---|---|---|---|
-| 0a | **`.stop_grad()` 是就地打标记（≠torch detach）** | 把它当 detach 用 → 原变量本体被标无梯度，前向全对、相关 loss 梯度静默为 0 | detach 语义一律用 `.detach()`（返回新变量）；`.stop_grad()` 只用于"这个变量从此不需要梯度"（如冻结参数） |
-| 0b | **任何 `for x in tensor` 式 python 循环建图 = O(N) 图节点** | 密集批次 grad/sync 阶段 100% CPU 停滞分钟级、位置随 shuffle 漂移、无报错。实锤案例：底座 `diag3d`（循环 jt.diag，gwd_loss 每 iter ×3）、逐实例 voronoi、逐实例 nonzero | 一律批量闭式改写（stack/where/searchsorted 分桶）；detach 路径转 numpy；用 `jt.liveness_info()` 打点确认 lived_ops 有界 |
-| 0c | **多进程 dataloader 环形缓冲死锁** | worker 卡 buffer.send、主进程卡 idqueue.pop（py-spy 可证） | num_workers=0（加载占比小时性价比最高） |
-| 1 | **jittor 1.3.8.5 + numpy≥2** | 凡 `jt.array`（numpy 数据）喂进计算图的算子输出未初始化内存；roundtrip 与无输入算子（rand/ones）正常 | 钉死 numpy==1.26.4；smoke 测试有防回归断言 |
-| 2 | **`jt.linalg.{det,inv,eigh,svd}` 是 numpy_code 实现** | GPU 上要 cupy（未装则 ModuleNotFoundError，且因惰性求值报错位置在千里之外的第一次 sync 处） | 2×2 全部用闭式解（`ops/linalg2x2.py`、`det_2x2`） |
-| 3 | **in-place 赋值的自动微分不可靠** | 前向对、梯度 0 | loss 路径一律 out-of-place（jt.where/concat 重建）；L2 必须比梯度 |
-| 4 | **`x.max()`/`.sum()` 等 reduce 返回 shape [1] 而非标量** | `jt.stack([a.max(), b.max()])` 得 (2,1)，下游 shape 错乱或 concat 报错 | 用 `jt.concat` 拼 reduce 结果 |
-| 5 | **GPU 的 `jt.matmul` 不广播 batch 维**（cublas_batched_matmul） | CPU 通过、GPU 报 Wrong inputs | batch 维不一致时显式 `expand`（见 solve_2x2） |
-| 6 | **conda 26.x `env config vars` 导出成大写** | jittor 读小写 `cc_path`，静默不生效 | 用 `etc/conda/activate.d/*.sh` |
-## 🟡 语义差异（有对应物但行为不同）
-| torch/mm 系 | jittor/JDet | 差异 |
-|---|---|---|
-| `torch.meshgrid(x, y, indexing='xy')` | 无 indexing 参数 | 手工构造：`X = x[None,:].expand((ny,nx))` |
-| mmengine `LinearParamScheduler` | — | **分母是 end-begin-1**（≠torch LinearLR），第 W-1 iter 到顶 |
-| mmcv `RoIAlignRotated(clockwise=True)` | JDet 原生方向 | **JDet 原生 = mmcv clockwise=True**（kernel 旋转矩阵互为转置）；aligned=-0.5 偏移需另加 |
-| mmcv `nms_rotated` keep 按分数降序 | jdet 返回原始顺序 | head 里先 argsort 再 NMS |
-| `torch.unique(return_inverse/counts)` | jt.unique 语义不保证 | detach 量一律走 numpy |
-| `Tensor.index_reduce_('amin'/'mean')` | 无 | amin 于 unique 分组=组内代表（scatter）；mean 带梯度 → one-hot 矩阵乘 |
-| `F.smooth_l1_loss(beta=)` | jt.nn.smooth_l1_loss 无 beta | 手写 `_smooth_l1` |
-| torchvision `resized_crop`（antialias） | jt.nn.interpolate 无 antialias | 已记录差异（增广路径） |
-| `nn.GroupNorm(requires_grad=)` | 无此参数 | config 适配层剥离该键 |
-| mmengine cfg 的 `axis=` 遗留 | jt.concat 只认 `dim=` | 底座 coder.py 已修 |
-| 底座 `RotatedResize` | mmdet Resize | 底座经 poly 往返会重规范角度（点框 0→-π/2）；自写 MMRotateResize |
-| 底座 `RotatedRandomFlip` | mmdet RandomFlip | 底座 x'=W-x-1 且角度 π-a；mmrotate 是 W-x 与 -a；自写 MMRotateRandomFlip |
-| 底座 `WhollyWoodDOTADataset` point_dummy=0.1 | 官方 ConvertWeakSupervision 默认 1 | v2 数据集显式用 1 |
-## 🟢 直接可用（已验证）
-`jt.concat/stack/split/gather/flip/clamp/where/arange/linspace/permute` /
-`nn.grid_sample(含 reflection)` / `nn.affine_grid` / `nn.interpolate(bilinear, align_corners)` /
-`nn.GroupNorm/PixelShuffle/ConvTranspose2d/MaxPool2d` / `jt.optim.AdamW(+clip_grad_norm 全局 L2)`。
-## 数值/梯度专题
-- **2×2 eigh 闭式解**（`ops/linalg2x2.py`）：`λ± = mean ± sqrt(clamp(delta²+b², 1e-24))`；
-  EPS 取 1e-24（1e-12 会污染面积 1e-6 的框）；b 对称化读取让梯度在 [0,1]/[1,0] 对称分摊
-  （对齐 torch.linalg.eigh 的梯度约定）；`|delta|/disc ≤ 1` ⇒ 特征值梯度天然有界，w==h 不炸。
-- **简并（特征值重复）时**特征基与逐特征值子梯度不唯一：与 torch 的差异只在 a/c 间分配，
-  trace 不变量一致。parity 测试对简并用例只比基无关量。
-- **PSC decode ±π/2 端点**：atan2(∓0,-1) 符号翻转给出 +π/2 vs -π/2，le90 周期 π 下等价，
-  按模 π 角距离比较。
-- **diff_iou_rotated 的鞋带公式必须在中心化坐标上��**（2026-07-26，stage-2 parity 揪出）：
-  mmcv 在原始图像坐标（~1e2-1e3）上算，x_i·y_j 项达 1e4-1e6，微小/退化交集依赖正负大项
-  相消；jittor 编译器把 `a*d - b*c` 融合成 FMA 后两项舍入不再互为相反数，残差 ~1e-2 px²
-  变成假面积（golden 上 loss_bbox 差 16%，torch eager 恰好精确抵消得 0）。修法：排序仍在
-  归一化坐标做（不变），面积改在「减均值后再乘 mask 归零填充槽」的坐标上算——闭合多边形
-  鞋带公式平移不变，且填充槽零贡献保持。torch 侧原始坐标的噪声底 ~1e-3 仍留在 golden 的
-  loss_bbox 梯度里，故 FCOS head 梯度紧验走 cls+ctr 支路（1e-4），总梯度 5e-3。
-- **jittor nn.GroupNorm 是一遍式方差** `E[x²]−E[x]²`（灾难性消去风险），torch 是两遍式
-  `E[(x−mean)²]`。底座 BRICKS 'GN' 已换成 `GroupNorm2Pass`（modules.py，签名/参数名兼容）。
-- **golden 若在 CUDA torch 上 dump，必须关 TF32**：torch 2.x 默认 `cudnn.allow_tf32=True`，
-  卷积只有 10 位尾数，前向自带 ~4e-4 相对误差，会整体污染 golden（曾被误判为 GN/实现差异，
-  实际逐层 CPU 对比 conv 输出 rel 1e-7）。dump 脚本统一加
-  `torch.backends.cudnn.allow_tf32 = False; torch.backends.cuda.matmul.allow_tf32 = False`。
-- **torch `ReLU(inplace=True)` 会篡改中间量取证**：dump 中间激活时先 `.clone()` 再过激活，
-  否则 pre-act 张量被原地覆盖，逐层对比出现「norm 层 rel=1.0 而 act 后 rel=1e-7」的假象。
-- **GPU 数值 parity 不可用于 bisect / 逐元素断言**（issue #3 结案，2026-07-26 22:00）：
-  jittor conv 前向/反向用 `cudnnFindConvolution*AlgorithmEx` **真实计时**选算法并做
-  **进程内 cache**；cudnn8 在 A100 上 `CUDNN_DEFAULT_MATH` 允许 TF32 tensor-core 算法
-  参与候选。选中 TF32 的进程 GPU-vs-golden 前向 ~4e-4、梯度 rel_l2 ~4e-2（逐元素违约
-  ~22%）；选中 FMA 的进程 ~1e-4。结果「进程内确定、跨进程漂移」，且**换 commit 改图形状
-  会扰动 benchmark 计时** → bisect 会把算法选择的翻转误报成某个 commit 的数值回归
-  （B 对 ba5b766 scatter-add 的 bisect 即此假信号；scatter vs one-hot 实测梯度逐位一致，
-  复现脚本 tools/debug/）。此前文档写的「共卡 cudnn 漂移、偶发 ~4%」是同一现象的
-  不完整解释。对策：梯度验收 = GPU 松验（rel_l2<5e-2，兜断链级 bug）+ CPU 紧验
-  （rel_l2<1e-4）；训练不受影响（TF32 卷积训练是业界默认，torch 侧同样如此）。
-## 工程
-- Jittor 首跑 JIT 编译数分钟属正常；变长 shape 反复触发重编译（R5）——M5 实测后决定 bucket。
-- 惰性求值：报错位置 ≠ 出错算子位置，二分法用 `.sync()`/`jt.flags.lazy_execution=0` 定位。
-- 权重转换：torch `state_dict` → `{k: numpy}` pickle → `model.load_parameters`（TED 已验证）。
----
-# v3 专项附录（Agent B）
-# Porting Notes（torch → Jittor）— 随移植进度持续更新
-## 2026-07-26 静态盘点（v3 核心四文件的 API 风险清单）
-统计对象：`point2rbox_v3_head.py`(H) / `point2rbox_v3.py`(D) /
-`point2rbox_v2_loss.py`(L, A 所有权) / `losses/utils.py`(U)。
-### 🔴 Jittor 无对应，必须自实现
-| API | 出现 | 对策 |
-|---|---|---|
-| `Tensor.index_reduce_(0, idx, src, 'amin'/'mean', include_self=False)` | H×9（实例聚合核心） | **已解**：底座 `h2rbox_v2p_head.py` L401-429 范式 = `np.unique`(CPU) 取 idx + Jittor 原生 `reindex_reduce('min'/'max'/'add')`（可导）；mean = add / count。照抄该范式即可，L1 测试仍要锁语义（include_self=False 等价性、空 segment） |
-| `torch.linalg.eigh` | D×1, L×5 | 全部 2×2 对称阵，闭式解（A 在 `ops/linalg2x2.py` 做，我验 `w==h` 退化） |
-| `torch.linalg.solve` | D×1, L×2 | 2×2 伴随矩阵求逆，det clamp eps |
-| `torch.diag_embed` | D×2, H×1, L×5 | jt.stack 手拼 2×2 或展开为标量表达式 |
-### 🟡 语义差异需 L1 锁死
-| API | 出现 | 风险点 |
-|---|---|---|
-| `torch.unique(return_inverse=True, return_counts=True)` | H×2 | jt.unique 的返回值/排序语义需实测；bid 聚合正确性完全依赖它 |
-| `grid_sample` / `affine_grid` | D×4 / D×2 | rotate_crop 用；align_corners 显式传；padding_mode='reflection' 行为需比对 |
-| `torch.gather(x, 1, idx)` | H×1 | jt.gather 语义核对 |
-| `topk` | H×2, L×7 | 排序稳定性、返回顺序 |
-| `F.interpolate bicubic` | D(voronoi val 上采样是 bilinear+align_corners=True) | align_corners 显式传 |
-| `cv2.watershed` / `cv2.medianBlur` | D, L | CPU numpy 路径，原样保留，确认梯度切断 |
-### 🔴 in-place 就地赋值（R5：前向对、梯度错）
-- H×4、D×4、L×1 处布尔掩码赋值，重点：
-  - `pos_decoded_angle_preds[square_mask] = 0`（H L303，**在梯度路径上**）
-  - `pos_rbox_targets[~pos_syn_mask, 2:] = pos_rbox_preds[...].detach()`���H L317，右侧已 detach，但左侧张量后续参与 loss_bbox）
-  - `data_sample.gt_instances.bboxes.tensor[mask] = results.bboxes.tensor`（D L417）
-- 对策：loss 路径一律改写为 `jt.where` / out-of-place 重组；L2 测试必须比梯度。
-### 🎲 随机性调用点（测试对齐时替换为固定序列）
-- D：`torch.rand`×5（dual-stream 分支选择、rot 角、scale 因子）+ `np.random.randint`（copy-paste 位置）
-- H：`torch.rand`×1
-- 测试策略：从 PyTorch dump 随机序列 npz，两边读同一序列（仅测试，训练不改）。
-### 来自 A 的 core 层结论（2026-07-26，M5' 必须遵守）
-1. **EdgeLoss 的 RoIAlignRotated**：必须用
-   `ROIAlignRotated(49, spatial_scale, sampling_ratio, aligned=True, clockwise=True)`
-   （A 在 commit 312b3b7 加的 mmcv 语义参数；**默认参数是底座旧行为，禁止用默认值**）。
-   实测 JDet 原生 kernel 方向 = mmcv clockwise=True；out_size 7/49 与 mmcv rtol 1e-4 全过。
-   残留差异：aligned=True 时 mmcv 不做 1×1 最小 roi 钳制、JDet kernel 始终钳制（亚像素 roi 才受影响）。
-2. **⚠️ nms_rotated 返回顺序**：jdet 的 keep 是**原始顺序**，mmcv 是**按分数降序**。
-   → v3 head 的 `_predict_by_feat_single`（test_cfg: nms 后 `max_per_img=2000` 截断）
-   **必须先按 score 降序 sort 再截断**，否则截掉的是错误的框，mAP 静默劣化。
-3. **LR 语义**：用 A 的 `LinearWarmupMultiStepLR`（与 mmengine 1440 点逐点相等）；
-   注意 mmengine LinearLR 分母是 `end-begin-1`（第 499 iter 到顶），且 A 已把
-   scheduler.step 挪到 optimizer.step 之前（f(i) 对齐）。自己写任何调度相关代码别按 torch 直觉来。
-4. GDLoss（A 修了 3 处底座语义差异）/ PSCCoder / box_iou_rotated 均已 parity 过，直接复用。
-5. **core-stable（13:29）随附三条**：(a) jt 的 `.max()`/`.sum()` reduce 出 shape [1] 非标量，
-   `jt.stack` 两个会得 (2,1)——聚合标量用 `jt.concat`；(b) jittor 1.3.8.5 无
-   `meshgrid(indexing=...)`，手工 `x[None,:].expand(...)` 构造；(c) w==h 简并时 eigh
-   逐特征值梯度是子梯度（基底相关，与 torch 不逐位一致但 trace 一致）——parity 测试对
-   简并用例只比不变量（trace/行列式/loss 值），不比逐元素梯度。
-6. `segment_anything()` 已在 A 的 loss 里 lazy import `jdet.models.sam`（我的签名）✓。
-   `RotatedSingleRoIExtractor` 落在 `models/roi_extractors/rotated_single_level.py`
-   （兼容 mmcv 旧参名 out_size/sample_num）。
-7. targets/bids 的 JDet 侧约定：待 A 移植 v2 detector（其 M4）时定，COORD 通知后照抄。
-### 已确认的良性点
-- head 只 build `PSCCoder` + `DistanceAnglePointCoder`，无 assigner（铁律三已证实）
-- MobileSAM 无 `nn.MultiheadAttention`（自定义 Attention，q/k/v 独立 Linear）→ 无 in_proj 拆包
-- MobileSAM 的 timm 依赖仅 DropPath/to_2tuple/trunc_normal_，推理期均可平替
-- TED 纯卷积，无自定义算子；`PixelShuffle(1)` 恒等（staging 已做防御性回退）
-### epoch 开关清单（SetEpochInfoHook 注入，L2 测试点）
-- `copy_paste_start_epoch=6`（D）、`edge_loss_start_epoch=6`（H L414）、
-  `joint_angle_start_epoch=1`（H L298：之前 angle detach）、
-  `label_assign_pseudo_label_switch_eopch=6`（D L403：切换 predict vs generate_pseudo_targets）
-## ⚠️ 勘误（2026-07-26 16:40，来自 A 的实战 + B 审计确认）
-**PLAN §8 的『`.detach()` → `.stop_grad()`』建议是错的。** jittor 的 `.stop_grad()`
-是**就地打标记并返回自身**（非 torch detach 语义）：把还要参与 loss 的变量本体标成
-无梯度 → 前向全对、部分 loss 梯度静默变 0。要 detach 语义用 `.detach()`（返回新变量）
-或 `.clone().stop_grad()`。仅对确定无其他梯度消费者的节点（转 numpy 前、buffer、
-新建切片节点）可直接 stop_grad。B 侧 21 处使用已逐一审计安全（loss 梯度 parity
-1.1e-5 佐证）；A 侧修复 10 处误用（commit 1aa8ad8）。

+# PyTorch to Jittor porting notes
+This document records the implementation choices that are required for numeric
+and gradient parity with the Point2RBox-v3 PyTorch reference.
+## MobileSAM
+The native Jittor implementation is in `python/jdet/models/sam/`:
+| File | Component |
+|---|---|
+| `tiny_vit.py` | TinyViT image encoder, MBConv blocks and window attention |
+| `prompt_encoder.py` | point/box/mask prompt encoding and random positional encoding |
+| `transformer.py` | two-way transformer |
+| `mask_decoder.py` | mask tokens, hypernetworks and IoU prediction |
+| `sam.py` | preprocessing and mask postprocessing |
+| `predictor.py` | `SamPredictor` and longest-side resize |
+| `build.py` | `sam_model_registry['vit_t']` and converted weight loading |
+The public API intentionally matches MobileSAM, so the loss code can use
+`sam_model_registry['vit_t']` and `SamPredictor` without an adapter layer.
+### Weight conversion
+PyTorch state dictionaries are converted to a pickle mapping parameter names
+to NumPy arrays:
+```bash
+python tools/convert_torch_weights.py mobile_sam.pt weights/mobile_sam.pkl
+```
+`num_batches_tracked` entries are ignored because Jittor BatchNorm does not use
+them. TinyViT `Conv2d_BN` keeps the original `c` and `bn` child names so the
+remaining keys load directly. The random positional-encoding matrix is a real
+checkpoint value, not a disposable initialization buffer; loading code and
+tests verify that it is restored.
+TinyViT caches indexed attention biases in evaluation mode. The cache is
+refreshed after loading weights, preventing stale random biases from surviving
+checkpoint restoration.
+### Tensor semantics
+- PyTorch `transpose(a, b)` swaps two dimensions. Jittor's transpose API is
+  expressed as a full permutation in this port; all window partition/reverse
+  paths therefore use explicit `permute` orders.
+- Point masking is written with `where` and mask multiplication instead of
+  in-place indexed assignment, preserving gradients.
+- `repeat_interleave` in the mask decoder is implemented with explicit
+  reshape/expand/reshape operations.
+- Predictor resizing uses PIL rather than torchvision and keeps the reference
+  longest-side coordinate transform.
+The converted 439-tensor checkpoint reaches mask IoU 1.0000 against PyTorch in
+CPU fp32 and at least 0.9982 on GPU in the repository tests.
+## TED edge detector
+TED is implemented in `python/jdet/models/edge/`. Its convolutional weights are
+converted with:
+```bash
+python tools/convert_ted_weights.py ted.pth weights/ted.pkl
+```
+The output edge maps match the PyTorch implementation with maximum relative
+error below 5.7e-6 in the validated CPU path.
+## Jittor semantic differences
+### Detach and in-place updates
+`Var.stop_grad()` changes the variable itself; it is not equivalent to
+PyTorch's `detach()`. Graph branches that require detached values use
+`Var.detach()`. Loss-path mask updates are rebuilt out-of-place with `where` or
+concatenation because indexed in-place writes can silently alter gradients.
+### Batched linear algebra
+Jittor's GPU `matmul` does not broadcast batch dimensions. Batch dimensions are
+expanded explicitly before batched products. The covariance operations used by
+Point2RBox are 2×2, so `python/jdet/ops/linalg2x2.py` provides closed-form
+determinant, solve and symmetric eigendecomposition implementations without a
+CuPy dependency.
+Repeated eigenvalues do not define a unique eigenbasis. Degenerate tests compare
+basis-independent quantities such as trace, determinant and reconstructed
+matrices rather than individual eigenvectors.
+### Reductions and grouping
+Some Jittor reductions return shape `(1,)` rather than a scalar. Scalar lists
+are concatenated instead of stacked. Instance grouping uses scatter reductions;
+Python loops over tensors are avoided because they create one graph fragment per
+element and can make dense batches stall.
+### Normalization
+Jittor's default GroupNorm computes variance as `E[x^2] - E[x]^2`. The port uses
+the two-pass implementation in `python/jdet/models/utils/modules.py` to match
+PyTorch and avoid cancellation. The detector explicitly redispatches
+`backbone.train()` after Jittor's recursive mode switch so frozen stages and
+`norm_eval=True` remain effective.
+### Rotated geometry
+- JDet's rotated RoIAlign angle direction matches mmcv
+  `clockwise=True`; Point2RBox uses `aligned=True` explicitly.
+- JDet rotated NMS returns retained indices in input order. Prediction code sorts
+  by score before truncating to `max_per_img`.
+- Rotated IoU polygon area is evaluated in centered coordinates. The shoelace
+  formula is translation invariant, while centering avoids false areas caused by
+  cancellation of large image-coordinate products.
+- `torchvision.transforms.functional.resized_crop(..., antialias=True)` is
+  reproduced by explicit antialiased bilinear weights in the scale augmentation
+  path.
+## Numeric validation
+The repository stores fixed PyTorch goldens under `tests/parity/golden/`.
+CPU fp32 checks use strict tolerances for losses and feature gradients. GPU
+checks use looser aggregate tolerances because cuDNN algorithm selection and
+TF32 can vary between processes even when the implementation is unchanged.
+Config values, including the class-specific SAM mask filtering table, are
+checked separately with zero tolerance.