Refactor: Move BB regressor debug tensor comparison to model_comparison/bbreg_comparison.py; modularize test/compare_models.py

test/compare_models.py

@@ -25,6 +25,8 @@ print("PYTHON SCRIPT: Set cuDNN benchmark=False, deterministic=True")
 from pytracking.features.net_wrappers import DiMPTorchScriptWrapper
 # For loading AtomIoUNet from source
 from ltr.models.bbreg.atom_iou_net import AtomIoUNet
+# Add import for new modular comparison
+from model_comparison.bbreg_comparison import compare_debug_tensors
 
 SCRIPT_DIR_FOR_INIT = os.path.dirname(os.path.abspath(__file__))
 ROOT_DIR_FOR_INIT = os.path.dirname(SCRIPT_DIR_FOR_INIT)
@@ -593,7 +595,7 @@ class ComparisonRunner:
                 try:
                     # Use from-source get_iou_feat for consistent 256-channel features
                     # DiMPTorchScriptWrapper.bb_regressor.get_iou_feat returns features with different channel counts
-                    temp_iou_feat = self.bb_regressor_from_source.get_iou_feat([py_feat_layer2, py_feat_layer3])
+                    temp_iou_feat = self.bb_regressor_from_source.get_iou_feat([py_feat_layer2, py_feat_layer3], i)
                     if isinstance(temp_iou_feat, tuple): temp_iou_feat = list(temp_iou_feat)
                     if len(temp_iou_feat) >= 2:
                         py_iou_feat_list = [temp_iou_feat[0], temp_iou_feat[1]]
@@ -1768,4 +1770,9 @@ if __name__ == "__main__":
     runner.run_all_tests()
     # runner.generate_html_report() # This is called within run_all_tests
     # print(f"HTML report generated at {runner.comparison_dir / 'report.html'}") # Old, caused TypeError
-    print(f"HTML report generated at {os.path.join(runner.comparison_dir, 'report.html')}") # New, fixed TypeError 
+    print(f"HTML report generated at {os.path.join(runner.comparison_dir, 'report.html')}") # New, fixed TypeError 
+    # Compare BB regressor debug tensors (BatchNorm and ReLU outputs)
+    cpp_bbreg_dir = os.path.join("test", "output", "bb_regressor")
+    py_bbreg_dir = os.path.join("test", "output_py", "bb_regressor")
+    print("\n=== BB Regressor Debug Tensor Comparison (BatchNorm and ReLU) ===")
+    compare_debug_tensors(cpp_bbreg_dir, py_bbreg_dir, sample_idx=0, verbose=True)

test/model_comparison/__init__.py

@@ -0,0 +1 @@
+ 

test/model_comparison/bbreg_comparison.py

@@ -0,0 +1,54 @@
+import os
+import torch
+import numpy as np
+
+def compare_debug_tensors(cpp_dir, py_dir, sample_idx=0, verbose=True):
+    """
+    Compare C++ and Python debug tensors for BB regressor (BatchNorm and ReLU outputs).
+    Args:
+        cpp_dir (str): Directory with C++ debug tensors.
+        py_dir (str): Directory with Python debug tensors.
+        sample_idx (int): Sample index to compare.
+        verbose (bool): Print detailed comparison results.
+    Returns:
+        dict: Comparison metrics for each layer/output.
+    """
+    layers = ["conv3_1t", "conv3_2t", "conv4_1t", "conv4_2t"]
+    stages = ["bn", "relu"]
+    results = {}
+    for layer in layers:
+        for stage in stages:
+            cpp_file = os.path.join(cpp_dir, f"sample_{sample_idx}_debug_{layer}_{stage}.pt")
+            py_file = os.path.join(py_dir, f"sample_{sample_idx}_debug_{layer}_{stage}_py.pt")
+            if not os.path.exists(cpp_file) or not os.path.exists(py_file):
+                results[f"{layer}_{stage}"] = None
+                if verbose:
+                    print(f"[WARN] Missing file: {cpp_file if not os.path.exists(cpp_file) else py_file}")
+                continue
+            cpp_tensor = torch.load(cpp_file, map_location="cpu")
+            py_tensor = torch.load(py_file, map_location="cpu")
+            cpp_np = cpp_tensor.cpu().detach().numpy().astype(np.float32)
+            py_np = py_tensor.cpu().detach().numpy().astype(np.float32)
+            if cpp_np.shape != py_np.shape:
+                results[f"{layer}_{stage}"] = None
+                if verbose:
+                    print(f"[WARN] Shape mismatch for {layer}_{stage}: C++ {cpp_np.shape}, Py {py_np.shape}")
+                continue
+            flat_cpp = cpp_np.flatten()
+            flat_py = py_np.flatten()
+            mae = np.mean(np.abs(flat_cpp - flat_py))
+            max_err = np.max(np.abs(flat_cpp - flat_py))
+            l2_cpp = np.linalg.norm(flat_cpp)
+            l2_py = np.linalg.norm(flat_py)
+            cos_sim = np.dot(flat_cpp, flat_py) / (l2_cpp * l2_py) if l2_cpp > 0 and l2_py > 0 else float('nan')
+            results[f"{layer}_{stage}"] = {
+                "mae": mae,
+                "max_err": max_err,
+                "cos_sim": cos_sim,
+                "l2_cpp": l2_cpp,
+                "l2_py": l2_py,
+                "shape": cpp_np.shape
+            }
+            if verbose:
+                print(f"{layer}_{stage}: cos_sim={cos_sim:.6f}, mae={mae:.4e}, max_err={max_err:.4e}, shape={cpp_np.shape}")
+    return results