File size: 3,922 Bytes
d6d6039 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | #!/usr/bin/env python3
"""
det03 — Convert det_500m.onnx to PyTorch (onnx2torch) and test batch independence.
Runs three forward passes with the same real image:
A: ONNXRuntime, B=1
B: onnx2torch, B=1
C: onnx2torch, B=5 (same image duplicated)
Prints C output shapes, then compares:
A vs B — confirms onnx2torch faithfully reproduces the ONNX computation
B vs C0 — checks if frame[0] of B=5 matches the B=1 output (batch independence test)
C0 vs C1 — checks if both frames in B=5 match each other (same input → same output)
"""
import argparse
import cv2
import numpy as np
import torch
import onnx2torch
import onnxruntime as ort
DET_SIZE = (640, 640)
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to det_500m.onnx')
parser.add_argument('--image', required=True, help='Path to any face image (jpg/png)')
args = parser.parse_args()
# ── Load models ───────────────────────────────────────────────────────────────
print('Loading ONNXRuntime session ...')
sess = ort.InferenceSession(args.model, providers=['CPUExecutionProvider'])
inp_name = sess.get_inputs()[0].name
print('Loading onnx2torch model ...')
pt_model = onnx2torch.convert(args.model)
pt_model.eval()
print('Done.\n')
# ── Preprocess image ──────────────────────────────────────────────────────────
img = cv2.imread(args.image)
assert img is not None, f'Could not read image: {args.image}'
print(f'Image: {args.image} original shape={img.shape}')
blob = cv2.dnn.blobFromImage(
cv2.resize(img, DET_SIZE),
1.0 / 128.0, DET_SIZE,
(127.5, 127.5, 127.5), swapRB=True,
) # shape: (1, 3, 640, 640)
t = torch.from_numpy(blob)
# ── A) ONNXRuntime B=1 ───────────────────────────────────────────────────────
ort_out = sess.run(None, {inp_name: blob})
# ── B) onnx2torch B=1 ────────────────────────────────────────────────────────
with torch.no_grad():
pt_out_b1 = pt_model(t)
# ── C) onnx2torch B=5 [img, img, img, img, img] ──────────────────────────────
t2 = torch.cat([t] * 5, dim=0)
with torch.no_grad():
pt_out_b2 = pt_model(t2)
# ── Print B=5 output shapes ───────────────────────────────────────────────────
print('\n[C] Output shapes for B=5:')
for i, c in enumerate(pt_out_b2):
print(f' out[{i}]: {list(c.shape)}')
# ── Compare ───────────────────────────────────────────────────────────────────
print('\n[D] Comparison across all 9 output heads\n')
print(f' {"head":<6} {"AvsB":>10} {"BvsC0":>10} {"C0vsC1":>10} note')
print(' ' + '-' * 60)
for i, (a, b, c) in enumerate(zip(ort_out, pt_out_b1, pt_out_b2)):
a = a.flatten()
b = b.detach().numpy().flatten()
c = c.detach().flatten()
K = b.shape[0]
c0 = c[:K]
c1 = c[K:]
diff_ab = np.abs(a - b).max()
diff_b_c0 = (torch.from_numpy(b) - c0).abs().max().item()
diff_c0_c1 = (c0 - c1).abs().max().item()
note = ''
if diff_b_c0 > 1e-4:
note = '<-- CONTAMINATION'
print(f' out[{i}] {diff_ab:>10.5f} {diff_b_c0:>10.5f} {diff_c0_c1:>10.5f} {note}')
print('\nPass criteria: BvsC0 < 1e-4 and C0vsC1 < 1e-4 for all heads.')
print('Done.')
|