GlobalNet for 3D affine registration won't train on a dataset where LocalNet works fine

[thierleft]

Hi @wyli
I have been running a training of LocalNet on 256x256x64 intensity images (intensity range 0.0-1.0) paired with their binary segmentation masks for deformable registration, computing and combining intensity-based loss (GlobalMutualInformation) and segmentation loss (GeneralizedDiceLoss). It works really well after 1500 epochs, but I would like to get affine transformations on this dataset and I have been trying (unsuccessfully) to train GlobalNet instead for this purpose. It seems that whatever I do, inputting intensity images only (with intensity-based loss), inputting segmentations only (with segmentation-based loss), numerous preprocessing methods on this dataset, tweaking network's hyperparameters and weights initialization, it won't train and the losses are never minimized (either stochastically moving around a fixed value no matter how many epochs or staying on the same loss value). Have you encountered similar cases when training GlobalNet for 3D affine registration or does anything come to mind?

[wyli]

maybe it's necessary to somehow include regularisers for the theta or the output

MONAI/monai/networks/nets/regunet.py

Lines 299 to 302 in 7a760e6

	if self.save_theta:
	self.theta = theta.detach()
	out: torch.Tensor = self.affine_transform(theta) - self.grid
	return out

@ebrahimebrahim probably have more ideas on this topic...

[ebrahimebrahim]

I have always struggled to get large affine transformations to work right using roughly the same approach you are taking, whether it's using GlobalNet or whether it's my own setup just using the built in monai UNet.

It's counterintuitive, but for some reason it's much easier to get a CNN to learn nonparameteric deformable registration than it is to get it to learn affine registration. I'd love to hear an explanation for why that is if anyone understands it.

I still haven't understood the problem of pairwise affine registration with a CNN, but here are some lessons I've learned trying to get it to work:

• When the affine transformations needed to align images in the dataset are small, things are better. It's possible to test this by starting with a dataset for which images are already all affine-aligned, e.g. to a population template, and then including random affine transforms as a data augmentation step to train something like GlobalNet. As you increase the magnitude of the transforms in the data augmentation, it becomes harder and harder to get a successful training session.
• Learning iterative transformations can often get you a little further. Here's what I mean: Let $M$ denote the moving image and let $F$ denote the fixed image. Normally we want to input $(M,F)$ into our model $\phi$ and we want to get out an affine transformation $A = \phi(M,F)$ such that $M.A$ is very similar to $F$. Instead of expecting the model output $A$ to completely align $M$ to $F$, we can demand that the output $A$ only gets $M$ a small step closer to alignment, such that repeating the process, say, 10 times would achieve complete alignment. So you could say $A_1 := \phi(M,F)$ and $M_1 := M.A_1$, then $A_2:=\phi(M_1,F)$ and $M_2 := M_1.A_2$, and so on until you get to $M_{10}$. Then your training objective would be based on similarity of $M_{10}$ to $F$, backpropping through all iterations.
  • Here is a work that uses the iterative approach: https://arxiv.org/abs/1903.08811 (section 2.1)
  • When actually implementing this you would want to compose transformations to avoid repeatedly resampling. So for example $M_2$ should be computed as $M.(A_1.A_2)$ rather than $(M.A_1).A_2$