1442 add LocalNet initialization

[kate-sann5100]

Signed-off-by: kate-sann5100 yiwen.li@st-annes.ox.ac.uk

Fixes #1442.

Description

Support different initialisations for LocalNet feature extraction layers.

Status

Ready/Work in progress/Hold

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 --codeformat --coverage.
• Quick tests passed locally by running ./runtests.sh --quick.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[yiheng-wang-nv]

Thanks @kate-sann5100 for the implementation. As for the mypy error, could you please modify the code as follow:

class LocalNetFeatureExtractorBlock(nn.Module):
    """
    A feature-extraction module that can be used for LocalNet, based on:
    `Weakly-supervised convolutional neural networks for multimodal image registration
    <https://doi.org/10.1016/j.media.2018.07.002>`_.
    `Label-driven weakly-supervised learning for multimodal deformable image registration
    <https://arxiv.org/abs/1711.01666>`_.

    Adapted from:
        DeepReg (https://github.com/DeepRegNet/DeepReg)
    """

    def __init__(
        self,
        spatial_dims: int,
        in_channels: int,
        out_channels: int,
        act: Optional[Union[Tuple, str]] = "RELU",
        initializer: str = "kaiming_uniform",
    ) -> None:
        """
        Args:
        spatial_dims: number of spatial dimensions.
        in_channels: number of input channels.
        out_channels: number of output channels.
        act: activation type and arguments. Defaults to ReLU.
        kernel_initializer: kernel initializer. Defaults to None.
        """
        super(LocalNetFeatureExtractorBlock, self).__init__()
        self.conv_block = get_conv_block(
            spatial_dims=spatial_dims, in_channels=in_channels, out_channels=out_channels, act=act, norm=None
        )
        conv_type: Type[Union[nn.Conv1d, nn.Conv2d, nn.Conv3d]] = Conv[Conv.CONV, spatial_dims]
        for m in self.conv_block.modules():
            if isinstance(m, conv_type):
                if initializer == "kaiming_uniform":
                    nn.init.kaiming_normal_(torch.as_tensor(m.weight))
                    print("kaiming")
                elif initializer == "zeros":
                    nn.init.zeros_(torch.as_tensor(m.weight))
                else:
                    raise ValueError(
                        f"initializer {initializer} is not supported, " "currently supporting kaiming_uniform and zeros"
                    )

And import Type and Conv as well:

from typing import Optional, Sequence, Tuple, Type, Union

import torch
from torch import nn
from torch.nn import functional as F

from monai.networks.blocks import Convolution
from monai.networks.layers import same_padding
from monai.networks.layers.factories import Conv, Norm, Pool

[kate-sann5100]

Thank you for the suggested fix, that's super helpful! @yiheng-wang-nv

[wyli]

/black

[wyli]

thanks!