Repository: Project-MONAI/MONAI

Issue:
return the spatial transform parameters in registration networks
**Is your feature request related to a problem? Please describe.**
follow up of https://github.com/Project-MONAI/MONAI/discussions/6128, the estimated spatial parameters should be accessible from the registration networks
https://github.com/Project-MONAI/MONAI/blob/e4b313d59b2702cd2a441d7a557650ced78dc6c8/monai/networks/nets/regunet.py#L251-L252

**Describe the solution you'd like**
set up an activation hook which is not intuitive to use
https://github.com/Project-MONAI/MONAI/discussions/6128#discussioncomment-5269348


**Additional context**
https://github.com/Project-MONAI/MONAI/discussions/6313 considers updating the metadata to include this piece of information.



Discussion:

Patch:
diff --git a/monai/networks/nets/regunet.py b/monai/networks/nets/regunet.py
--- a/monai/networks/nets/regunet.py
+++ b/monai/networks/nets/regunet.py
@@ -65,7 +65,7 @@ def __init__(
             out_activation: activation at the last layer
             out_channels: number of channels for the output
             extract_levels: list, which levels from net to extract. The maximum level must equal to ``depth``
-            pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv3d
+            pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv
             concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition
             encode_kernel_sizes: kernel size for down-sampling
         """
@@ -234,7 +234,22 @@ def forward(self, x):
 
 
 class AffineHead(nn.Module):
-    def __init__(self, spatial_dims: int, image_size: list[int], decode_size: list[int], in_channels: int):
+    def __init__(
+        self,
+        spatial_dims: int,
+        image_size: list[int],
+        decode_size: list[int],
+        in_channels: int,
+        save_theta: bool = False,
+    ):
+        """
+        Args:
+            spatial_dims: number of spatial dimensions
+            image_size: output spatial size
+            decode_size: input spatial size (two or three integers depending on ``spatial_dims``)
+            in_channels: number of input channels
+            save_theta: whether to save the theta matrix estimation
+        """
         super().__init__()
         self.spatial_dims = spatial_dims
         if spatial_dims == 2:
@@ -255,6 +270,9 @@ def __init__(self, spatial_dims: int, image_size: list[int], decode_size: list[i
         self.fc.weight.data.zero_()
         self.fc.bias.data.copy_(out_init)
 
+        self.save_theta = save_theta
+        self.theta = torch.Tensor()
+
     @staticmethod
     def get_reference_grid(image_size: tuple[int] | list[int]) -> torch.Tensor:
         mesh_points = [torch.arange(0, dim) for dim in image_size]
@@ -278,6 +296,8 @@ def forward(self, x: list[torch.Tensor], image_size: list[int]) -> torch.Tensor:
         f = x[0]
         self.grid = self.grid.to(device=f.device)
         theta = self.fc(f.reshape(f.shape[0], -1))
+        if self.save_theta:
+            self.theta = theta.detach()
         out: torch.Tensor = self.affine_transform(theta) - self.grid
         return out
 
@@ -305,7 +325,22 @@ def __init__(
         pooling: bool = True,
         concat_skip: bool = False,
         encode_kernel_sizes: int | list[int] = 3,
+        save_theta: bool = False,
     ):
+        """
+        Args:
+            image_size: output displacement field spatial size
+            spatial_dims: number of spatial dims
+            in_channels: number of input channels
+            num_channel_initial: number of initial channels
+            depth: input is at level 0, bottom is at level depth.
+            out_kernel_initializer: kernel initializer for the last layer
+            out_activation: activation at the last layer
+            pooling: for down-sampling, use non-parameterized pooling if true, otherwise use conv
+            concat_skip: when up-sampling, concatenate skipped tensor if true, otherwise use addition
+            encode_kernel_sizes: kernel size for down-sampling
+            save_theta: whether to save the theta matrix estimation
+        """
         for size in image_size:
             if size % (2**depth) != 0:
                 raise ValueError(
@@ -315,6 +350,7 @@ def __init__(
                 )
         self.image_size = image_size
         self.decode_size = [size // (2**depth) for size in image_size]
+        self.save_theta = save_theta
         super().__init__(
             spatial_dims=spatial_dims,
             in_channels=in_channels,
@@ -334,6 +370,7 @@ def build_output_block(self):
             image_size=self.image_size,
             decode_size=self.decode_size,
             in_channels=self.num_channels[-1],
+            save_theta=self.save_theta,
         )
 
 


Test patch:
diff --git a/tests/test_globalnet.py b/tests/test_globalnet.py
--- a/tests/test_globalnet.py
+++ b/tests/test_globalnet.py
@@ -21,11 +21,11 @@
 from monai.networks.blocks import Warp
 from monai.networks.nets import GlobalNet
 from monai.networks.nets.regunet import AffineHead
-from tests.utils import test_script_save
+from tests.utils import assert_allclose, test_script_save
 
 TEST_CASES_AFFINE_TRANSFORM = [
     [
-        {"spatial_dims": 3, "image_size": (2, 2, 2), "decode_size": (2, 2, 2), "in_channels": 1},
+        {"spatial_dims": 3, "image_size": (2, 2, 2), "decode_size": (2, 2, 2), "in_channels": 1, "save_theta": True},
         torch.ones(2, 12),
         torch.tensor([[[1, 2], [2, 3]], [[2, 3], [3, 4]]]).unsqueeze(0).unsqueeze(0).expand(2, 3, 2, 2, 2),
     ],
@@ -55,10 +55,12 @@
             "pooling": True,
             "concat_skip": True,
             "encode_kernel_sizes": 3,
+            "save_theta": theta,
         },
         (1, 1, 16, 16),
         (1, 2, 16, 16),
     ]
+    for theta in (False, True)
 ]
 
 
@@ -66,6 +68,8 @@ class TestAffineHead(unittest.TestCase):
     @parameterized.expand(TEST_CASES_AFFINE_TRANSFORM)
     def test_shape(self, input_param, theta, expected_val):
         layer = AffineHead(**input_param)
+        if input_param.get("save_theta"):
+            assert_allclose(layer.theta, torch.Tensor())
         result = layer.affine_transform(theta)
         np.testing.assert_allclose(result.cpu().numpy(), expected_val.cpu().numpy(), rtol=1e-4, atol=1e-4)
 
@@ -81,13 +85,15 @@ def test_shape(self, input_param, input_shape, expected_shape):
         with eval_mode(net):
             img = torch.randn(input_shape)
             result = net(img.to(device))
+            if input_param.get("save_theta"):
+                assert_allclose(net.output_block.theta, torch.Tensor([[1.0, 0.0, 0.0, 0.0, 1.0, 0.0]]))
             warped = warp_layer(img.to(device), result)
             self.assertEqual(result.shape, expected_shape)
             # testing initial pred identity
             np.testing.assert_allclose(warped.detach().cpu().numpy(), img.detach().cpu().numpy(), rtol=1e-4, atol=1e-4)
 
-    def test_script(self):
-        input_param, input_shape, _ = TEST_CASES_GLOBAL_NET[0]
+    @parameterized.expand(TEST_CASES_GLOBAL_NET)
+    def test_script(self, input_param, input_shape, _):
         net = GlobalNet(**input_param)
         test_data = torch.randn(input_shape)
         test_script_save(net, test_data)