How to change the data loaders?

[LIKP0]

Hello all,

Thanks for your great work first! I'm using nnUNetv2 for 3D brain MRI segmentation and the data shape is 192,160,224. Now I want to add one special sample data in my traning step with 50% probability like this:

def train_step(self, batch: dict) -> dict:
        data = batch['data'] # List: 5, [(2, 36, 128, 128, 128), (2, 36, 64, 64, 64), ..., (2, 36, 32, 32, 32)]
        target = batch['target']

        data = data.to(self.device, non_blocking=True)
        if isinstance(target, list):
            target = [i.to(self.device, non_blocking=True) for i in target]
        else:
            target = target.to(self.device, non_blocking=True)

        self.optimizer.zero_grad(set_to_none=True)
        with autocast(self.device.type, enabled=True) if self.device.type == 'cuda' else dummy_context():
            if probablity > 50%:
                output = self.network(data) # 50% use original data
            else
                output = self.network(my_special_sample) # 50% use special data
            l = self.loss(output, target)  # nnUNet base loss: compute loss on 5 scale and aggregate

            with torch.no_grad():
                ema_output = self.ema_model(self.get_noisy_input(data))  # only use top scale
            uncertainty = self.MC_uncertainty(self.ema_model, data)  # forward T times on ema_model to get uncertainty
            l2 = self.UA_consistency_loss(output[0], ema_output[0], uncertainty)
           ......

I find nnUNet does a downsampling for data so they become a list [(2, 36, 128, 128, 128), (2, 36, 64, 64, 64), ..., (2, 36, 32, 32, 32)], how can I also do this for my special sample data? This special data is just one file and I just want to simply load and use it.

The code of the dataloaders get me a little confused. (Yeah I admit my poor coding skills...) Could anyone give some advice? Thanks in advance!

[LIKP0]

Okay, I think maybe an extra dataset can help, i.e., to put the special sample as another dataset with nnUNet preprocessing. And just add that dataset to the input of train step like def train_step(self, batch: dict, batch2: dict) -> dict:

Does anyone have better ideas?

[Karol-G]

Dear @LIKP0,

Thank you for using nnU-Net and reaching out with your query.

To assist you effectively, it would be helpful if you could provide more detailed information regarding your issue. Specifically, could you clarify what exactly you are trying to achieve with the special sample data and how it relates to your overall training process? A more comprehensive understanding of your goal and the problem you're encountering will enable us to offer more precise guidance.

Looking forward to your response.

Best regards, Karol Gotkowski

[LIKP0]

Hello @Karol-G ,

Thank you for your attention.

I have completed my desired research with nnUNet, and I can share my experience here:

What i'm going to do is to combine the uncertainty-aware mean teacher (UAMT) framework with nnUNet. Like the figure below, I only need to change the UNet segmentation network to nnUNet.

The figure references to Xu Z, Lu D, Luo J, et al. Anti-interference from noisy labels: Mean-teacher-assisted confident learning for medical image segmentation[J]. IEEE Transactions on Medical Imaging, 2022, 41(11): 3062-3073. Fig.2

However, I need to provide two independent dataset low-quality (LQ) and high-quality (HQ) to student and teacher model seperately. Finally, I add an extra dataloader through the whole pipeline just by copying similar nnUNet code. Now the only problem is the code is too dirty hahhhhhh :-(

Maybe next step, the nnUNet team could embed nnUNet into frequently-used deep learning framework like mean teacher MT? Just a little advice :-)

Thanks for all your great work again!

Best wishes, JiaLin Li

[Karol-G]

Hey @Karol-G,

Sounds like a very cool project and it is nice to hear that you managed to solve it already yourself! Also thank you for your advice, but as it is in real life there is always too much one wants to do and too little time to actually do it ;)

Best regards, Karol