Adaptive refactorization

[mmuckley]

A draft for refactorizations for #205 to reduce duplicate code.

[Timsey]

On a first pass:

Possibly major stuff:

• EquispacedMaskFunc is gone from subsample.py and instead we have EquiSpacedMaskFunc and EquispacedMaskFractionFunc. EquispacedMaskFractionFunc seems to have the same functionality as EquispacedMaskFunc from the adaptive VarNet code, but it is missing a block controlled by the skip_low_freqs argument (we used this to create masks that had the proper number of lines when combined with dense low-frequency sampling): https://github.com/facebookresearch/fastMRI/blob/58c0bfcd676630cc3f002a8129a1388dc8896307/fastmri/data/subsample.py#L212-L233

Minor stuff:

• Comment here is now a lie and we can just remove it: https://github.com/facebookresearch/fastMRI/blob/69db5c4e49ed0eda93bae8ed4daefacd6b2235b5/fastmri/models/active_varnet.py#L145
• Argparse description for val_volume_sample_rate and test_volume_sample_rate should refer to val_sample_rate and test_sample_rate (and vice-versa), not to sample_rate: https://github.com/facebookresearch/fastMRI/blob/69db5c4e49ed0eda93bae8ed4daefacd6b2235b5/fastmri/pl_modules/data_module.py#L370
• Argument active_acquisition in MiniCoilTransform is unused and can be removed everywhere in that class: https://github.com/facebookresearch/fastMRI/blob/69db5c4e49ed0eda93bae8ed4daefacd6b2235b5/fastmri/data/transforms.py#L525
• Example folder is now called _activevarnet, but should be called _adaptivevarnet. Same for active_varnet.py and active_varnet_module.py.
• MiniCoilTransform and VarNetTransform now also return full kspace. This is necessary for sign-fixing in the adaptive VarNet, but not for the base VarNet, which does now also receive kspace in the forward(). Not necessarily bad, but I just thought I'd mention it as a potential thing to clean up.

[mmuckley]

@Timsey could you explain the first a bit more? Do you just want the acceleration mask where the densely-sampled center does not exist?

[Timsey]

@Timsey could you explain the first a bit more? Do you just want the acceleration mask where the densely-sampled center does not exist?

So the original EquispacedMaskFunc that I used (now EquispacedMaskFractionFunc) does not fully take into account the densely sampled center when computing accel_samples here: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri/data/subsample.py#L337 This does compute an adjusted acceleration, but it can still end up sampling one or more of its lines in the center, leading to overlapping samples and varying actual acceleration factors (based on e.g. offset).

The quoted code block ensures a constant acceleration by forcing the equisipaced sampling to happen outside of the low-frequency region. The downside of this is that samples on one side of the low-frequency region are often not exactly equispaced w.r.t. samples on the other side of the low-frequency region.

EDIT: In any case, the latter is what I used for the Equispaced experiments in the paper.

Second edit: Yeah, essentially this is just dense center mask + 'acceleration mask where the densely-sampled center does not exist', except for offset issues.

[mmuckley]

Sorry I should have commented here: @Timsey I added your version of the EquispacedMaskFunc to the fastmri_examples folder for your project.

[Timsey]

LGTM! One more thing I noticed is that with MriModule reverted to the original, a bunch of the stuff we log with wandb is missing now, so I would propose just removing that logging.

EDIT: Nevermind, this functionality is now in ActiveVarNetModule.

@luisenp We should probably also add that single-checkpoint evaluation script?

[mmuckley]

@luisenp @Timsey We do have single-checkpoint eval sketches here if it helps, but you all can use whatever format you prefer.

https://github.com/facebookresearch/fastMRI/blob/main/fastmri_examples/varnet/run_pretrained_varnet_inference.py

[Timsey]

Random comment: this line here can be simplified, and the comment removed: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri_examples/adaptive_varnet/subsample.py#L128. Simplification is just to remove the entire line and change skip here to num_low_frequencies: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri_examples/adaptive_varnet/subsample.py#L134.

[Timsey]

Another issue, this one a bit larger.

This bit of code will crash if we ever reach it: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri_examples/adaptive_varnet/train_adaptive_varnet_demo.py#L411-L415

because VarNetDataTransform returns a VarNetSample: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri/data/transforms.py#L500-L509

but AdaptiveVarNetModule expects a tuple as returned by MiniCoilTransform: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri/data/transforms.py#L646-L655

see: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri/pl_modules/adaptive_varnet_module.py#L197

This also creates an issue for training equispaced models with our training script, because VarNetModule expects a VarNetSample: https://github.com/facebookresearch/fastMRI/blob/mmuckley/adaptive_varnet_updates/fastmri/pl_modules/adaptive_varnet_module.py#L197

So we either have to: 1) Change MiniCoilTransform and AdaptiveVarNetModule to accept VarNetSample (note that this VarNetSample also needs to contain the full kspace for sign-fixing purposes). 2) Make our own copy VarNetDataTransform that returns a tuple like MiniCoilTransform, and our own copy of VarNetModule that expects a tuple like AdaptiveVarNetModule.

[Timsey]

So we either have to:

1. Change MiniCoilTransform and AdaptiveVarNetModule to accept VarNetSample (note that this VarNetSample also needs to contain the full kspace for sign-fixing purposes).
2. Make our own copy VarNetDataTransform that returns a tuple like MiniCoilTransform, and our own copy of VarNetModule that expects a tuple like AdaptiveVarNetModule.

I'm thinking option 2 is better now, since I'm also running into some issues with VarNetModule not expecting any of the later-added args, such as num_sense_lines.

[mmuckley]

@Timsey I can spend some time on these. How are you launching jobs locally?