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.
64 lines
2.2 KiB
64 lines
2.2 KiB
import torch
|
|
import os
|
|
from ltr.models.backbone import resnet50
|
|
|
|
def save_tensor(tensor, path):
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
torch.save(tensor.cpu(), path)
|
|
|
|
# Load the preprocessed input tensor exported from C++
|
|
preprocessed_input_path = 'test/input_samples/common/sample_0_image.pt'
|
|
input_obj = torch.load(preprocessed_input_path, map_location='cpu', weights_only=False)
|
|
if isinstance(input_obj, torch.Tensor):
|
|
input_tensor = input_obj
|
|
else:
|
|
# Try to extract tensor from JIT module
|
|
if hasattr(input_obj, 'named_parameters'):
|
|
params = list(input_obj.named_parameters())
|
|
if params:
|
|
input_tensor = params[0][1]
|
|
else:
|
|
raise RuntimeError('No tensor found in JIT module')
|
|
else:
|
|
raise RuntimeError('Unknown input object type')
|
|
|
|
# Add batch dimension if missing
|
|
if input_tensor.dim() == 3:
|
|
input_tensor = input_tensor.unsqueeze(0)
|
|
|
|
# Load Python ResNet model
|
|
model = resnet50(output_layers=['layer1', 'layer2', 'layer3', 'layer4'], pretrained=False)
|
|
# Optionally load weights if needed (uncomment and set path if required)
|
|
# model.load_state_dict(torch.load('backbone_pure_tensors/state_dict.pt'))
|
|
model.eval()
|
|
|
|
# Forward pass and save intermediate outputs
|
|
x = input_tensor
|
|
out_dir = 'test/output_py/resnet_debug/'
|
|
|
|
# After conv1
|
|
x1 = model.conv1(x)
|
|
save_tensor(x1[0], os.path.join(out_dir, 'sample_0_after_conv1.pt'))
|
|
# After bn1
|
|
x2 = model.bn1(x1)
|
|
save_tensor(x2[0], os.path.join(out_dir, 'sample_0_after_bn1.pt'))
|
|
# After relu1
|
|
x3 = model.relu(x2)
|
|
save_tensor(x3[0], os.path.join(out_dir, 'sample_0_after_relu1.pt'))
|
|
# After maxpool
|
|
x4 = model.maxpool(x3)
|
|
save_tensor(x4[0], os.path.join(out_dir, 'sample_0_after_maxpool.pt'))
|
|
# After layer1
|
|
x5 = model.layer1(x4)
|
|
save_tensor(x5[0], os.path.join(out_dir, 'sample_0_after_layer1.pt'))
|
|
# After layer2
|
|
x6 = model.layer2(x5)
|
|
save_tensor(x6[0], os.path.join(out_dir, 'sample_0_after_layer2.pt'))
|
|
# After layer3
|
|
x7 = model.layer3(x6)
|
|
save_tensor(x7[0], os.path.join(out_dir, 'sample_0_after_layer3.pt'))
|
|
# After layer4
|
|
x8 = model.layer4(x7)
|
|
save_tensor(x8[0], os.path.join(out_dir, 'sample_0_after_layer4.pt'))
|
|
|
|
print('Saved all intermediate outputs for sample 0 using C++-preprocessed input.')
|