You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
70 lines
2.9 KiB
70 lines
2.9 KiB
import os
|
|
import torch
|
|
import numpy as np
|
|
|
|
def extract_tensor_from_jit(obj):
|
|
# Try to extract a tensor from a TorchScript module
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj
|
|
# Try .tensor attribute (used in some containers)
|
|
if hasattr(obj, 'tensor') and isinstance(obj.tensor, torch.Tensor):
|
|
return obj.tensor
|
|
# Try named_parameters
|
|
if hasattr(obj, 'named_parameters'):
|
|
params = list(obj.named_parameters())
|
|
if params:
|
|
return params[0][1]
|
|
# Try named_buffers
|
|
if hasattr(obj, 'named_buffers'):
|
|
bufs = list(obj.named_buffers())
|
|
if bufs:
|
|
return bufs[0][1]
|
|
return None
|
|
|
|
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")
|
|
cpp_tensor = None
|
|
py_tensor = None
|
|
if os.path.exists(cpp_file):
|
|
try:
|
|
obj = torch.load(cpp_file, map_location="cpu", weights_only=False)
|
|
cpp_tensor = extract_tensor_from_jit(obj)
|
|
except Exception as e:
|
|
print(f"[WARN] Could not load {cpp_file}: {e}")
|
|
else:
|
|
print(f"[WARN] Missing file: {cpp_file}")
|
|
if os.path.exists(py_file):
|
|
try:
|
|
obj = torch.load(py_file, map_location="cpu", weights_only=False)
|
|
py_tensor = extract_tensor_from_jit(obj)
|
|
except Exception as e:
|
|
print(f"[WARN] Could not load {py_file}: {e}")
|
|
else:
|
|
print(f"[WARN] Missing file: {py_file}")
|
|
if cpp_tensor is None or py_tensor is None:
|
|
print(f"Warning: Cannot compare 'BBReg Debug_{layer.capitalize()}' for sample {sample_idx}, one or both tensors are None.")
|
|
continue
|
|
# Flatten and compare
|
|
cpp_flat = cpp_tensor.detach().cpu().numpy().flatten()
|
|
py_flat = py_tensor.detach().cpu().numpy().flatten()
|
|
cos_sim = np.dot(cpp_flat, py_flat) / (np.linalg.norm(cpp_flat) * np.linalg.norm(py_flat) + 1e-12)
|
|
results[(layer, stage)] = cos_sim
|
|
if verbose:
|
|
print(f"Cosine similarity for {layer} {stage}: {cos_sim:.6f}")
|
|
return results
|