Implemented VoxelMorph

[kvttt]

Fixes #5484.

Description

Implemented VoxelMorph and added some docstrings. Checked coding style locally.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[wyli]

/black

CCing @ebrahimebrahim @brudfors @nvahmadi in case you have comments/suggestions here..

[wyli]

thanks, it looks good to me in general, have you used this to train some baseline successfully? also please help make the test cases smaller if possible (like input size (2, 1, 96, 96, 48)?), so they can run with minimal system resources of the github runner.

[nvahmadi]

Thanks for tagging me, @wyli and congrats @kvttt - this looks like a valuable contribution, awesome in fact! Thanks a lot! :) Some thoughts on this PR:

• I like the input parameter for integration steps and that a number int_steps>0 automatically switches on diffeomorphic registration. To be more explicit, I would recommend to rename the parameter to "integration_steps", because the prefix "int" reminds me of "integer" and the suffix "steps" alone does not really explain what this parameter is about.
• Also, I saw that you implemented the construction of a U-Net all from scratch. I wonder whether there was an explicit reason to reimplement this? We have the U-Net class, the DynUNet class, both are nicely parametrizable, and the VoxelMorph backbone can be mimicked with these classes (e.g. by setting an appropriate number of channels etc). An advantage over the reimplementation would be that the existing MONAI (Dyn)UNet implementations offer a lot of extra functionality (e.g. in-layer residuals, anisotropic kernels and strides etc.), which all would need to be reimplemented in this class, if we wanted to cover them.
• Thinking this a bit further, the VoxelMorph approach does not even require a U-Net backbone, it can be any fully convolutional architecture that maps the input to a same-sized 3-channel output (either DVF or DDF). For example, in a recent tutorial, we used a SegResNet as the mapping backbone. One could even consider letting the user choose the backbone at constructor level (either as a string, or as a class definition), and letting appropriate network params be passed as an additional constructor parameter (a dictionary, in this case). The default class would then be a UNet without in-layer residuals and the params set up in this PR (to mimic the original VoxelMorph paper). Overall, this would allow much more flexibility for the user to explore different backbones, e.g. UNet, DynUNet, AttentionUNet, SegResNet and event 3D-ViTs like (Swin-)UNETR etc.

Apart form that though, as said, I think this is a very welcome addition to MONAI Core. I would recommend to also create an accompanying tutorial. We could use the OASIS subtask from the Learn2Reg challenge as a dataset. Feel free to DM me to sync on this, we are working on this as well, would be great to collaborate on this to avoid duplicate efforts - ideally we could even use your new VoxelMorph class. :)

[brudfors]

Thanks @kvttt!

I agree with @nvahmadi that it would be preferable to make use of the networks that are already in MONAI, so that there are not multiple implementations of a UNet let's say. Looking at the Paired Lung CT 3d Registration tutorial; another way of including a VoxelMorph model in MONAI could be to simply add a new tutorial to MONAI, similar to the Lung CT 3d Registration, but replacing the LocalNet with a UNet (and whatever hyper-parameters that were used in the VoxelMorph paper) and switching its forward function with the one implemented in @kvttt's PR. Perhaps with the OASIS brain scans used in the learn2reg challenge, as @nvahmadi proposed.

[wyli]

in the current implementation there's already self.net decoupling the backbone from the rest of the network, perhaps we can further make self.net configurable by providing an argument backbone

my_convnet = MyConvNet(...)
model = VoxelMorph(backbone=my_convnet, ...)

[kvttt]

thanks, it looks good to me in general, have you used this to train some baseline successfully? also please help make the test cases smaller if possible (like input size (2, 1, 96, 96, 48)?), so they can run with minimal system resources of the github runner.

Thanks, @wyli for the instructions. I have not tried running a baseline yet. But I can set up one right now using the Learn2Reg OASIS subtask as @nvahmadi mentioned.

As a response to @nvahmadi's second question: I did not use the existing monai.networks.nets.UNet because the _get_down_layer method does not support using maxpooling (used by default in VoxelMorph) instead of strided convolution for downsampling.

MONAI/monai/networks/nets/unet.py

Lines 200 to 204 in c3f9914

	def _get_down_layer(self, in_channels: int, out_channels: int, strides: int, is_top: bool) -> nn.Module:
	"""
	Returns the encoding (down) part of a layer of the network. This typically will downsample data at some point
	in its structure. Its output is used as input to the next layer down and is concatenated with output from the
	next layer to form the input for the decode (up) part of the layer.

And _get_up_layer by default uses transpose convolution for upsampling which is not used in VoxelMorph.

MONAI/monai/networks/nets/unet.py

Lines 275 to 277 in c3f9914

	conv_only=is_top and self.num_res_units == 0,
	is_transposed=True,
	adn_ordering=self.adn_ordering,

I have not looked at monai.networks.nets.DynUNet yet but I could try building the UNet portion of VoxelMorph from it.

in the current implementation there's already self.net decoupling the backbone from the rest of the network, perhaps we can further make self.net configurable by providing an argument backbone

my_convnet = MyConvNet(...)
model = VoxelMorph(backbone=my_convnet, ...)

I think this probably would be easier to implement and more intuitive to use, where one is given the option to specify which backbone they want. One can then opt to use the vanilla VoxelMorph or specify any other networks as suggested by @nvahmadi.

[nvahmadi]

Thanks @kvttt, good point, the UNets in MONAI are all based on convolution for down-/upsampling, whereas VoxelMorph was still relying on pooling. It's a good idea to keep your implementation as the "original VoxelMorph", but offering to pass a different backbone, as @wyli suggested. Great to hear that you're supportive of this idea. :) This will be a very nice and flexible addition.

[ebrahimebrahim]

This looks amazing! And I agree with earlier comments about have an "original voxelmorph mode" and a "provide your own backbone" mode, including decoupling the current construction of the self.net into its own VoxelMorphConvNet (or whatever name makes sense). Continuing from the comment in #7178 (comment) one could then build the original voxelmorph with

my_convnet = VoxelMorphConvNet(in_channels, ...)
model = VoxelMorph(backbone=my_convnet, ...`) 

[kvttt]

In response to @wyli's comment

in the current implementation there's already self.net decoupling the backbone from the rest of the network, perhaps we can further make self.net configurable by providing an argument backbone

my_convnet = MyConvNet(...)
model = VoxelMorph(backbone=my_convnet, ...)

and @ebrahimebrahim's comment,

This looks amazing! And I agree with earlier comments about have an "original voxelmorph mode" and a "provide your own backbone" mode, including decoupling the current construction of the self.net into its own VoxelMorphConvNet (or whatever name makes sense). Continuing from the comment in #7178 (comment) one could then build the original voxelmorph with

my_convnet = VoxelMorphConvNet(in_channels, ...)
model = VoxelMorph(backbone=my_convnet, ...`) 

I made the distinction clear in my code between VoxelMorph (referring to the framework) versus VoxelMorphUNet (the backbone network) and added the option to specify a backbone network.

[wyli]

/build

[wyli]

/build

[wyli]

/build