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.
159 lines
7.3 KiB
159 lines
7.3 KiB
import os
|
|
import torch
|
|
import numpy as np
|
|
|
|
def extract_tensor_from_jit(obj):
|
|
"""Extract a tensor from a TorchScript RecursiveScriptModule or return the tensor if already a tensor."""
|
|
import torch
|
|
if isinstance(obj, torch.Tensor):
|
|
return obj
|
|
# Try common attribute names
|
|
for attr in ["tensor", "weight", "value"]:
|
|
if hasattr(obj, attr):
|
|
t = getattr(obj, attr)
|
|
if isinstance(t, torch.Tensor):
|
|
return t
|
|
# Try named_parameters and named_buffers
|
|
if hasattr(obj, "named_parameters"):
|
|
for name, param in obj.named_parameters():
|
|
if isinstance(param, torch.Tensor):
|
|
return param
|
|
if hasattr(obj, "named_buffers"):
|
|
for name, buf in obj.named_buffers():
|
|
if isinstance(buf, torch.Tensor):
|
|
return buf
|
|
# Try children recursively
|
|
if hasattr(obj, "children"):
|
|
for child in obj.children():
|
|
t = extract_tensor_from_jit(child)
|
|
if t is not None:
|
|
return t
|
|
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")
|
|
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")
|
|
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
|
|
|
|
def compare_resnet_debug_tensors(cpp_dir, py_dir, sample_idx=0, verbose=True):
|
|
"""
|
|
Compare C++ and Python debug tensors for ResNet (after conv1, bn1, relu, maxpool, layer1, layer2, layer3, layer4).
|
|
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 stage.
|
|
"""
|
|
# Map stage to (cpp_filename, py_filename)
|
|
stage_map = {
|
|
"after_conv1": (f"sample_{sample_idx}_conv1_output.pt", f"sample_{sample_idx}_after_conv1.pt"),
|
|
"after_bn1": (f"sample_{sample_idx}_bn1_output.pt", f"sample_{sample_idx}_after_bn1.pt"),
|
|
"after_relu1": (f"sample_{sample_idx}_relu1_output.pt", f"sample_{sample_idx}_after_relu1.pt"),
|
|
"after_maxpool": (f"sample_{sample_idx}_maxpool_output.pt", f"sample_{sample_idx}_after_maxpool.pt"),
|
|
"after_layer1": (f"sample_{sample_idx}_layer1.pt", f"sample_{sample_idx}_after_layer1.pt"),
|
|
"after_layer2": (f"sample_{sample_idx}_layer2.pt", f"sample_{sample_idx}_after_layer2.pt"),
|
|
"after_layer3": (f"sample_{sample_idx}_layer3.pt", f"sample_{sample_idx}_after_layer3.pt"),
|
|
"after_layer4": (f"sample_{sample_idx}_layer4.pt", f"sample_{sample_idx}_after_layer4.pt"),
|
|
}
|
|
results = {}
|
|
import torch, os
|
|
for stage, (cpp_file, py_file) in stage_map.items():
|
|
cpp_path = os.path.join(cpp_dir, cpp_file)
|
|
py_path = os.path.join(py_dir, py_file)
|
|
cpp_tensor = None
|
|
py_tensor = None
|
|
if os.path.exists(cpp_path):
|
|
cpp_obj = torch.load(cpp_path, map_location='cpu', weights_only=False)
|
|
cpp_tensor = extract_tensor_from_jit(cpp_obj)
|
|
else:
|
|
if verbose:
|
|
print(f"[WARN] Missing file: {cpp_path}")
|
|
if os.path.exists(py_path):
|
|
py_obj = torch.load(py_path, map_location='cpu')
|
|
py_tensor = extract_tensor_from_jit(py_obj)
|
|
else:
|
|
if verbose:
|
|
print(f"[WARN] Missing file: {py_path}")
|
|
if cpp_tensor is None or py_tensor is None:
|
|
results[stage] = {
|
|
'cosine_similarity': None,
|
|
'allclose': None,
|
|
'max_abs_diff': None,
|
|
'cpp_shape': None if cpp_tensor is None else tuple(cpp_tensor.shape),
|
|
'py_shape': None if py_tensor is None else tuple(py_tensor.shape),
|
|
}
|
|
if verbose:
|
|
print(f"Warning: Cannot compare 'ResNet Debug {stage}' for sample {sample_idx}, one or both tensors are None.")
|
|
continue
|
|
# Flatten for cosine similarity
|
|
cpp_flat = cpp_tensor.flatten().float()
|
|
py_flat = py_tensor.flatten().float()
|
|
if cpp_flat.shape != py_flat.shape:
|
|
results[stage] = {
|
|
'cosine_similarity': None,
|
|
'allclose': False,
|
|
'max_abs_diff': None,
|
|
'cpp_shape': tuple(cpp_tensor.shape),
|
|
'py_shape': tuple(py_tensor.shape),
|
|
}
|
|
if verbose:
|
|
print(f"[WARN] Shape mismatch for {stage}: C++ {cpp_tensor.shape}, Py {py_tensor.shape}")
|
|
continue
|
|
cos_sim = torch.nn.functional.cosine_similarity(cpp_flat, py_flat, dim=0, eps=1e-8).item()
|
|
allclose = torch.allclose(cpp_tensor, py_tensor, atol=1e-5, rtol=1e-3)
|
|
max_abs_diff = (cpp_tensor - py_tensor).abs().max().item()
|
|
results[stage] = {
|
|
'cosine_similarity': cos_sim,
|
|
'allclose': allclose,
|
|
'max_abs_diff': max_abs_diff,
|
|
'cpp_shape': tuple(cpp_tensor.shape),
|
|
'py_shape': tuple(py_tensor.shape),
|
|
}
|
|
if verbose:
|
|
print(f"ResNet Debug {stage}: cos_sim={cos_sim:.6f}, allclose={allclose}, max_abs_diff={max_abs_diff:.3e}, shape={cpp_tensor.shape}")
|
|
return results
|