Ablation study, reproducibility and determinism

[Luugaaa]

Context

The current model with the following configuration produces very promising results.
Overall Configuration (not exhaustive) :

• region based training with minimal number of classes
• multichannel with phase + mag
• initial LR = 0.001
• LR scheduler = Cosine
• Optimizer = AdamW
• Mag preprocessing : light clahe
• Phase preprocessing : heavy contrast enhancement
• Augmentation : light spatial augmentation

Most of those hyperparameters were chosen based on intuition and personal experience. We'd like to make ablation studies to understand which hyperparameters really improve the results. We would also like to investigate new possible source of improvement such as (not exhaustive) :

• multichannel with mag & phase and with adjacent slices
• stronger augmentations
• stronger/weaker/no preprocessing

Goal

To make it possible to draw conclusions from experiments with different hyperparameters, we would like to make the results as reproducible as possible. The best way to do that is to enable deterministic training.

Issue

nnUnet does not currently support deterministic training, see issue 1423 of the nnunet repo. And no evolution is planned regarding that matter.

Current investigation

I added the basics that should make the training deterministic, but it's not yet.

seed=42
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
# I'm on mps so no CUDA-specific seeding is needed
torch.use_deterministic_algorithms(True)

I then noticed that even the dataloaders were not deterministic, so changed this code, to this :

        if allowed_num_processes == 0:
            mt_gen_train = SingleThreadedAugmenter(dl_tr, None)
            mt_gen_val = SingleThreadedAugmenter(dl_val, None)
        else:
            train_seeds = [MASTER_SEED + i for i in range(allowed_num_processes)]
            
            num_val_processes = max(1, allowed_num_processes // 2)
            val_seeds = [MASTER_SEED + 1000 + i for i in range(num_val_processes)] # Use an offset to avoid overlap

            mt_gen_train = MultiThreadedAugmenter(data_loader=dl_tr, transform=None,
                                                num_processes=allowed_num_processes,
                                                num_cached_per_queue=max(6, allowed_num_processes // 2), 
                                                seeds=train_seeds,
                                                pin_memory=self.device.type == 'cuda', wait_time=0.002)
            
            mt_gen_val = MultiThreadedAugmenter(data_loader=dl_val, transform=None, 
                                                num_processes=num_val_processes,
                                                num_cached_per_queue=max(3, allowed_num_processes // 4),
                                                seeds=val_seeds,
                                                pin_memory=self.device.type == 'cuda',
                                                wait_time=0.002)

Now the dataloaders are deterministic, but the training still isn't.

Next step

• Remove augmentation, this may be the source of non determinism

Possible workaround

I'm currently training 10 models (6/10 right now) on 40 epochs to have an estimate of the variability of the loss and metrics during training. If we cannot make nnunet deterministic, we can use this to have a better understanding of the impact of each hyperparameters. Note : this will involves training (on 40 epochs) several times for each hyperparameters modification, this is very time consuming and not ideal

[Luugaaa]

The random data augmentations (custom or default) are indeed the source of the non determinism ! When there are removed we have the following results, two trainings' (green & purple) metrics/loss overlap perfectly, the training is now deterministic !

[Luugaaa]

Even when we override the augmentation code in the custom trainer (where we seed everything), the augmentation still breaks the determinism.

Investigation of the determinism of the default augmentations

The elastic deform from the Spatial augmentation is a source of determinism breaking. The source of the non determinism is this code. The issue is the mixing of randomness from numpy and torch. The author does useless back and forth between torch and numpy, most likely because of how this method was originally implemented. It can be fixed (if you're willing to make changes in this module) by replacing it with this code :

            offsets_np = np.random.randn(dim, *self.patch_size).astype(np.float32)

            # all the additional time elastic deform takes is spent here
            for d in range(dim):
                # Perform FFT, filtering, and IFFT entirely in NumPy
                tmp = np.fft.fftn(offsets_np[d])
                tmp = fourier_gaussian(tmp, sigmas[d])
                offsets_np[d] = np.fft.ifftn(tmp).real

                mx = np.max(np.abs(offsets_np[d]))
                offsets_np[d] /= (mx / np.clip(magnitude[d], a_min=1e-8, a_max=np.inf))
            
            # Convert the final NumPy array to a PyTorch tensor
            spatial_dims = tuple(list(range(1, dim + 1)))
            offsets = torch.from_numpy(offsets_np)
            offsets = torch.permute(offsets, (*spatial_dims, 0))

[Luugaaa]

It seems to be a common error made again and again.

In the random transform (that encapsulates most augmentations), replace this line with this code :

return {"apply_transform": np.random.rand() < self.apply_probability}

For the gaussian noise, replace this line with this code :

dct['apply_to_channel'] = np.random.rand(shape[0]) < self.p_per_channel

and those lines with this code :

        if not isinstance(params['sigmas'], list):
            num_channels = sum(params['apply_to_channel'])
            # Generate noise with NumPy
            noise_np = np.random.normal(0, params['sigmas'], size=(1, *img_shape[1:]))
            # Convert to torch tensor
            gaussian = torch.from_numpy(noise_np.astype(np.float32))
            gaussian = gaussian.expand((num_channels, *[-1]*(len(img_shape) - 1)))
        else:
            # Generate noise with NumPy
            noise_np_list = [
                np.random.normal(0, i, size=(1, *img_shape[1:])) for i in params['sigmas']
            ]
            # Convert to torch tensors and concatenate
            gaussian = torch.cat(
                [torch.from_numpy(n.astype(np.float32)) for n in noise_np_list],
                dim=0
            )
        return gaussian

[Luugaaa]

For the gaussian blur, replace this line with this :

dct['apply_to_channel'] = np.random.rand(shape[0]) < self.p_per_channel

And make sure to use benchmark=False !

For the MultiplicativeBrightnessTransform replace this line with this :

        apply_to_channel_np = np.where(np.random.rand(shape[0]) < self.p_per_channel)[0]
        apply_to_channel = torch.from_numpy(apply_to_channel_np)

The BrightnessAdditiveTransform is already good.

In the ContrastTransform your need to change this line to :

    apply_to_channel_np = np.where(np.random.rand(shape[0]) < self.p_per_channel)[0]
    apply_to_channel = torch.from_numpy(apply_to_channel_np)

In SimulateLowResolutionTransform, import numpy as np and change this into :

    if self.allowed_channels is None:
        apply_to_channel_np = np.where(np.random.rand(shape[0]) < self.p_per_channel)[0]
        apply_to_channel = torch.from_numpy(apply_to_channel_np)
    else:
        apply_to_channel = [i for i in self.allowed_channels if np.random.rand() < self.p_per_channel]

Now in GammaTransform, import numpy as np and change this into that :

    apply_to_channel_np = np.where(np.random.rand(shape[0]) < self.p_per_channel)[0]
    apply_to_channel = torch.from_numpy(apply_to_channel_np)
    
    retain_stats = torch.from_numpy(np.random.rand(len(apply_to_channel)) < self.p_retain_stats)
    invert_image = torch.from_numpy(np.random.rand(len(apply_to_channel)) < self.p_invert_image)

And finally in MirrorTransform, import numpy as np and change this into that :

axes = [i for i in self.allowed_axes if np.random.rand() < 0.5]

[Luugaaa]

Now you've got yourself a nnunet training fully deterministic with the default augmentation pipeline !!

You should, of course, check if similar changes are needed when using other transforms.

[Luugaaa]

@jcohenadad I spent the day on this nnunet deterministic investigation, I'm thinking this might be interesting for SCT, as reproducibility was mentioned (I think) in a previous meeting.

This is also a nnunet issue that has no answer, so I might as well give one to the people who were asking, and propose a PR both in nnunet and in batchgeneratorsv2 (same author, big source of non determinism) if you think it's worth it !

[jcohenadad]

@jcohenadad I spent the day on this nnunet deterministic investigation, I'm thinking this might be interesting for SCT, as reproducibility was mentioned (I think) in a previous meeting.

Very cool investigation! And yes, it might be useful for the team. I suggest you mention it during our ivadomed and SCT meetings.

This is also a nnunet issue that has no answer, so I might as well give one to the people who were asking, and propose a PR both in nnunet and in batchgeneratorsv2 (same author, big source of non determinism) if you think it's worth it !

Yes! Definitely

[Luugaaa]

PRs submitted : here and here. We'll see what they do with it !