85 lines
3.1 KiB
Markdown
85 lines
3.1 KiB
Markdown
Work only on the Hiera backbone port and its stage-by-stage parity.
|
||
|
||
Goal:
|
||
Implement an MLX Hiera backbone that reproduces the feature contract expected by CorridorKey’s downstream heads.
|
||
|
||
Context:
|
||
|
||
- The original PyTorch model uses a timm backbone created with features_only=True.
|
||
- The downstream model expects a pyramid of 4 multiscale feature maps.
|
||
- The first input projection is patched from 3 input channels to 4 input channels for RGB + alpha hint.
|
||
- This phase is about reproducing backbone behavior and feature outputs, not full end-to-end inference.
|
||
|
||
Primary deliverables:
|
||
|
||
- src/corridorkey_mlx/model/hiera.py
|
||
- tests/test_hiera_stage_shapes.py
|
||
- tests/test_hiera_stage_parity.py
|
||
- updates to converter mapping files as needed
|
||
- README notes describing the backbone parity status and any known gaps
|
||
|
||
Requirements:
|
||
|
||
1. Inspect the original CorridorKey backbone usage and the underlying Hiera/timm feature contract before editing.
|
||
2. Treat the PyTorch reference harness as the source of truth for:
|
||
- number of feature stages
|
||
- output order
|
||
- channel counts
|
||
- spatial reductions
|
||
3. Implement the smallest MLX Hiera subset necessary to produce the correct 4 feature maps.
|
||
4. Preserve the patched 4-channel patch-embed behavior exactly:
|
||
- do not invent a new initialization rule
|
||
- match the PyTorch conversion semantics for the extra alpha-hint channel
|
||
5. Keep tensor layout handling explicit and centralized:
|
||
- use one canonical boundary between PyTorch-style NCHW fixtures and MLX NHWC internals
|
||
- do not scatter transpose logic across the backbone code
|
||
6. Add narrow parity checks for:
|
||
- patch embed output
|
||
- each stage output
|
||
- final list of multiscale features
|
||
7. If exact numerical parity is not immediately possible, first achieve:
|
||
- correct stage count
|
||
- correct stage order
|
||
- correct shapes
|
||
- correct dtype flow
|
||
then localize the first divergent block
|
||
8. Keep the implementation modular:
|
||
- patch embed
|
||
- stage/block definition
|
||
- downsampling / pooling transitions
|
||
- feature collection
|
||
9. Update the checkpoint conversion path only as needed to support the backbone weights for completed modules.
|
||
|
||
Diagnostics to produce:
|
||
|
||
- stage name
|
||
- source tensor shape
|
||
- destination tensor shape
|
||
- layout assumption
|
||
- max abs error
|
||
- mean abs error
|
||
- whether the mismatch begins before or after a stage transition
|
||
|
||
Do not:
|
||
|
||
- wire the full GreenFormer end-to-end yet
|
||
- rewrite decoder/refiner code unless required by an interface mismatch
|
||
- optimize with compile() yet
|
||
- add training code
|
||
- silently accept stage reordering or “close enough” feature contracts
|
||
|
||
Working style:
|
||
|
||
- Explore first.
|
||
- Summarize the exact Hiera feature contract you believe CorridorKey expects before editing.
|
||
- After edits, run the narrowest backbone-only tests first.
|
||
- When parity fails, stop at the first divergent stage and explain the likely cause before changing more code.
|
||
|
||
Definition of done:
|
||
|
||
- MLX backbone returns 4 multiscale feature maps in the correct order
|
||
- shape tests pass
|
||
- stage-level parity tests exist and run
|
||
- the 4-channel patch embed path is implemented and documented
|
||
- converter mappings exist for completed backbone weights
|