sub batching volumetric images to subvolumes for segmentation

[Jason-Fung]

Hello,

I'm new to the pytorch lightning API and its community, and I've been trying to figure out a way to refactor my original code to train a unet using MONAI for 3D volumetric image segmentation. One problem I encountered was figuring out a way to control the training of batches of subvolumes from a full image to avoid OOM issues. Looking at the pseudocode for DataParallel from this link: https://pytorch-lightning.readthedocs.io/en/stable/common/lightning_module.html#training-with-dataparallel
I thought it would be a good idea to override the tbptt_split_batch function with my own implementation.

def tbptt_split_batch(self, batch):

    upper,upper_shape,lower,lower_shape,full_mask,upper_mask,lower_mask = batch

    upper_batch = []
    lower_batch = []

    # shuffle dataset

    if self.shuffle == True:
            # shuffle the batches
            upper_key_list = list(range(len(upper))) # batch[0] is the upper volume
            random.shuffle(upper_key_list)
            
            # check if lower img exists, otherwise perform shuffling
            if len(lower) == 0:
                pass
            else:
                lower_key_list = list(range(len(lower)))
                random.shuffle(lower_key_list)
    else:
        upper_key_list = list(range(len(upper))) #
        lower_key_list = list(range(len(lower)))

    # split upper volume
    for i in range(0, len(upper), self.batch_size):
        if i + split_size >= len(upper):
            batch_keys = upper_key_list[i:len(upper)]

            sub_imgs = torch.squeeze(torch.stack([upper.get(key) for key in batch_keys], dim=1), dim = 0) 
            sub_masks = torch.squeeze(torch.stack([upper_mask.get(key) for key in batch_keys], dim=1), dim = 0)
            upper_batch.append(["upper",batch_keys,sub_imgs,sub_masks])
        else:
            batch_keys = upper_key_list[i:i+self.batch_size]

            sub_imgs = torch.squeeze(torch.stack([upper.get(key) for key in batch_keys], dim=1), dim = 0) 
            sub_masks = torch.squeeze(torch.stack([upper_mask.get(key) for key in batch_keys], dim=1), dim = 0)
            upper_batch.append(["upper",batch_keys, sub_imgs, sub_masks])

    # split lower volume
    for i in range(0, len(lower), self.batch_size):
        if i + split_size >= len(lower):
            batch_keys = lower_key_list[i:len(lower)]

            sub_imgs = torch.squeeze(torch.stack([lower.get(key) for key in batch_keys], dim=1), dim = 0) 
            sub_masks = torch.squeeze(torch.stack([lower_mask.get(key) for key in batch_keys], dim=1), dim = 0)
            lower_batch.append(["lower",batch_keys, sub_imgs, sub_masks])
        else:
            batch_keys = upper_key_list[i:i+self.batch_size]

            sub_imgs = torch.squeeze(torch.stack([lower.get(key) for key in batch_keys], dim=1), dim = 0) 
            sub_masks = torch.squeeze(torch.stack([lower_mask.get(key) for key in batch_keys], dim=1), dim = 0)
            lower_batch.append(["lower",batch_keys, sub_imgs, sub_masks])

    full_mask_list = ["full_mask",upper_shape, lower_shape, full_mask]

    return upper_batch + lower_batch + full_mask_list

Since my images have varying with even/odd stack sizes i.e., 34, 71, 85, etc. the strategy is to split the image into an upper volume (z_given % z_upper_volume == 0) and lower volumes (remainder stacks) then patching the upper volume and lower volumes into 3d patches i.e., (16,160,160).

I'm just wondering if this is the correct function to override if I want to do subvolume batch training?