updated collapse logic

[jamilgafur]

Merge branch 'fix/architecture-collapse-routing'

Fix module replacement routing and PyPrune LCA hook capture for complex architectures

This merge resolves critical routing and trace failures during the network discovery and collapse phases. Previously, the framework struggled with complex multi-branch models (InceptionNet, XceptionNet) and interleaved architectures (ConvNeXt), resulting in orphaned submodules and failed activation hooks.

Changes included:

1. Block Replacement Engine (collapse.py)

• Fix destructive patching: Fixed a bug in _build_and_replace_block where complex containers were blindly overwritten by standard nn.Sequential blocks. We now explicitly pass container=info.get("container") down to _replace_layers.
• Preserve custom forward logic: Enabled in-place patching (via setattr) for localized multi-branch containers, ensuring custom forward() routing inside models like InceptionNet and XceptionNet remains intact during surgery.
• Enhanced Diagnostics: Upgraded logging across the discovery phase to include strict spatial shape tracking, pre/post parameter delta reporting, and robust try/except wrappers around downstream validation.

2. ConvNeXt Architecture (ConvNetX.py)

• Align registration with execution: Refactored stage and downsample layer registration to perfectly mirror the computational graph flow.
• Flattened containers: Replaced disjointed self.stages and self.downsample_layers ModuleLists with a single, interleaved self.features nn.Sequential block.
• Static Graphing: Removed dynamic for loops from forward() and _calculate_fc_input_size(). This ensures PyPrune's LCA (Lowest Common Ancestor) locator can successfully trace and hook subgraphs that bridge across downsample and stage boundaries without throwing out-of-order execution errors.

[Copilot]

Pull request overview

This PR updates PyPrune’s “collapse”/surgery pipeline to better handle complex model topologies (multi-branch + manually routed containers) by improving LCA-based block localization, activation capture, and non-destructive in-place container patching. It also refactors the ConvNeXt model registration to provide a flatter, trace-friendly module graph and updates transfer candidate acceptance criteria to allow either parameter or FLOP reductions.

Changes:

• Refactors collapse.py to (a) resolve block boundaries via LCA, (b) capture both pre-block inputs and post-block outputs via hooks, and (c) patch non-nn.Sequential containers in-place to preserve custom forward() routing.
• Refactors ConvNetX.py to register ConvNeXt stages/downsamples as one continuous self.features container and removes forward loops.
• Updates transfer region selection logic to accept candidates when either params or FLOPs improve, and adds new run outputs / summary artifacts under manuscript/Tranfer/.

Reviewed changes

Copilot reviewed 21 out of 128 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
pyPrune/models/ConvNetX.py	Flattens stage/downsample registration into a single self.features sequential to better match execution order for tracing/collapse.
manuscript/Tranfer/collapse.py	Implements LCA-based container resolution, I/O hook capture, in-place container patching, and expanded diagnostics/validation steps.
manuscript/Tranfer/transfer.py	Adjusts dynamic candidate acceptance to allow either params or FLOPs reduction; updates combined-region validation accordingly.
manuscript/Tranfer/ConvNeXt_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for ConvNeXt CIFAR-10 (control-only).
manuscript/Tranfer/InceptionNet_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for InceptionNet CIFAR-10.
manuscript/Tranfer/MobileNet_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for MobileNet CIFAR-10 (control-only).
manuscript/Tranfer/RegNetX_400MF_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for RegNetX_400MF CIFAR-10.
manuscript/Tranfer/VGG16_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for VGG16 CIFAR-10.
manuscript/Tranfer/XceptionNet_Cifar10_epochs3_pretrain3_JF_discovered_regions.json	Adds discovered-region output for XceptionNet CIFAR-10.
manuscript/Tranfer/VGG16_Cifar10_discover_JF_False_pre3_ep3_run.txt	Adds run log artifact for VGG16 CIFAR-10 discovery workflow.
manuscript/Tranfer/MobileNet_Cifar10_discover_JF_False_pre3_ep3_run.txt	Adds run log artifact for MobileNet CIFAR-10 discovery workflow.
manuscript/Tranfer/master_results_summary_ep3_pre3.json	Adds aggregated results summary JSON for the ep3/pre3 experiment set.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/VGG16_cifar100__all_candidates.tex	Adds LaTeX table artifact for VGG16 CIFAR-100 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/VGG16_cifar100__all_candidates.csv	Adds CSV artifact for VGG16 CIFAR-100 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/RegNetX_400MF_cifar100__all_candidates.tex	Adds LaTeX table artifact for RegNetX_400MF CIFAR-100 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/RegNetX_400MF_cifar100__all_candidates.csv	Adds CSV artifact for RegNetX_400MF CIFAR-100 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/RegNetX_400MF_cifar10__all_candidates.tex	Adds LaTeX table artifact for RegNetX_400MF CIFAR-10 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/RegNetX_400MF_cifar10__all_candidates.csv	Adds CSV artifact for RegNetX_400MF CIFAR-10 candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/InceptionNet_tinyimagenet_all_candidates.tex	Adds LaTeX table artifact for InceptionNet TinyImageNet candidate blocks.
manuscript/Tranfer/figures_ep3_pre3/hardware_efficiency/InceptionNet_tinyimagenet_all_candidates.csv	Adds CSV artifact for InceptionNet TinyImageNet candidate blocks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.