Stage 1: Dump Hiera backbone outputs.
(predictor, img_batch)
| 65 | |
| 66 | |
| 67 | def dump_backbone(predictor, img_batch): |
| 68 | """Stage 1: Dump Hiera backbone outputs.""" |
| 69 | print("\n=== Stage 1: Hiera Backbone ===") |
| 70 | model = predictor.model |
| 71 | |
| 72 | # Run backbone trunk |
| 73 | with torch.no_grad(): |
| 74 | backbone_out = model.image_encoder.trunk(img_batch) |
| 75 | |
| 76 | # backbone_out is a list of intermediate features |
| 77 | # For Hiera, the output is typically from the trunk's forward which returns intermediates |
| 78 | print(f" Backbone output type: {type(backbone_out)}") |
| 79 | |
| 80 | if isinstance(backbone_out, (list, tuple)): |
| 81 | for i, feat in enumerate(backbone_out): |
| 82 | save_tensor(f"hiera_stage_{i}", feat) |
| 83 | print(f" hiera_stage_{i}: {feat.shape}") |
| 84 | elif isinstance(backbone_out, dict): |
| 85 | for k, v in backbone_out.items(): |
| 86 | if isinstance(v, torch.Tensor): |
| 87 | save_tensor(f"backbone_{k}", v) |
| 88 | print(f" backbone_{k}: {v.shape}") |
| 89 | |
| 90 | return backbone_out |
| 91 | |
| 92 | |
| 93 | def dump_fpn_neck(predictor, img_batch): |