| import unittest |
|
|
| import numpy as np |
| import onnxruntime as ort |
| import torch |
|
|
| from tools.maya_export.common import ckpt_paths, work_dir |
| from tools.maya_export.streamed_backbone import _Weights, shard_files |
|
|
| BUILD = work_dir() / "depth_build" / "fp32" |
| REF = work_dir() / "golden" / "ref" / "conv_01_s0.npz" |
| FRAMES = 20 |
| N_LEVELS = 15 |
|
|
|
|
| def _sess(path): |
| so = ort.SessionOptions() |
| so.intra_op_num_threads = 4 |
| return ort.InferenceSession(str(path), so, providers=["CPUExecutionProvider"]) |
|
|
|
|
| class TestDepthLevelParity(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| if not all((BUILD / f"level_{i:02d}.onnx").exists() for i in range(1, N_LEVELS + 1)): |
| raise unittest.SkipTest("per-level depth graphs not built") |
| if not REF.exists(): |
| raise unittest.SkipTest("reference corpus not found") |
| cls.weights = _Weights(shard_files(ckpt_paths()["ckpt"])) |
|
|
| def test_chain_matches_reference_depth_logits(self): |
| """The reference is stored in fp16. Logits here reach about 23, where one fp16 ulp is 0.0156, so the |
| check is: fp16-rounded outputs equal the stored values except for one-ulp rounding ties (<= 1%), and the |
| absolute error is at most half an ulp. The fp32 port is checked against upstream PyTorch fp32 separately |
| (tools/maya_export/depth_check.py).""" |
| with np.load(REF) as z: |
| codes = z["codes"][:, :FRAMES].T.astype(np.int64) |
| h = z["h_frames"][:FRAMES].astype(np.float32) |
| ref = z["depth_logits"][:FRAMES].astype(np.float32) |
| mem = self.weights.get("depth.null_mem").float().numpy().reshape(1, 1, 1024) |
| mem = np.repeat(mem, FRAMES, axis=0) |
| mismatch, total, worst = 0, 0, 0.0 |
| for level in range(1, N_LEVELS + 1): |
| sess = _sess(BUILD / f"level_{level:02d}.onnx") |
| got = sess.run(None, {"h": h, "codes": codes, "mem": mem})[0] |
| del sess |
| want = ref[:, level - 1] |
| g16 = got.astype(np.float16).astype(np.float32) |
| mismatch += int((g16 != want).sum()) |
| total += want.size |
| worst = max(worst, float(np.abs(g16 - want).max())) |
| print(f"fp16-rounded mismatches {mismatch}/{total}, max abs {worst:.7f}", flush=True) |
| self.assertLessEqual(mismatch / total, 0.01) |
| self.assertLessEqual(worst, 0.0078125) |
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
