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.
45 lines
1.8 KiB
45 lines
1.8 KiB
import torch
|
|
import os
|
|
from ltr.models.bbreg.atom_iou_net import AtomIoUNet
|
|
|
|
def compare_tensors(a, b, label):
|
|
a = a.float().cpu().contiguous().view(-1)
|
|
b = b.float().cpu().contiguous().view(-1)
|
|
if a.shape != b.shape:
|
|
print(f"{label}: Shape mismatch: {a.shape} vs {b.shape}")
|
|
return
|
|
cos_sim = torch.nn.functional.cosine_similarity(a, b, dim=0).item()
|
|
mae = torch.mean(torch.abs(a - b)).item()
|
|
max_abs = torch.max(torch.abs(a - b)).item()
|
|
print(f"{label}: cos_sim={cos_sim:.8f}, MAE={mae:.8e}, max_abs={max_abs:.8e}")
|
|
|
|
def main():
|
|
# Layer names and file patterns
|
|
layers = [
|
|
("conv3_1t", 0),
|
|
("conv3_2t", 0),
|
|
("conv4_1t", 0),
|
|
("conv4_2t", 0),
|
|
]
|
|
model_dir = "exported_weights/bb_regressor"
|
|
# Load Python model
|
|
model = AtomIoUNet(input_dim=(512, 1024), pred_input_dim=(256, 256), pred_inter_dim=(256, 256))
|
|
# Load weights from exported files
|
|
state_dict = {}
|
|
for lname, idx in layers:
|
|
wfile = os.path.join(model_dir, f"{lname}_{idx}_weight.pt")
|
|
bfile = os.path.join(model_dir, f"{lname}_{idx}_bias.pt")
|
|
state_dict[f"{lname}.{idx}.weight"] = torch.load(wfile, map_location='cpu', weights_only=False)
|
|
state_dict[f"{lname}.{idx}.bias"] = torch.load(bfile, map_location='cpu', weights_only=False)
|
|
model.load_state_dict(state_dict, strict=False)
|
|
# Compare each weight and bias
|
|
for lname, idx in layers:
|
|
py_w = getattr(getattr(model, lname)[idx], 'weight').data
|
|
py_b = getattr(getattr(model, lname)[idx], 'bias').data
|
|
cpp_w = state_dict[f"{lname}.{idx}.weight"]
|
|
cpp_b = state_dict[f"{lname}.{idx}.bias"]
|
|
compare_tensors(py_w, cpp_w, f"{lname}_weight")
|
|
compare_tensors(py_b, cpp_b, f"{lname}_bias")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|