Training code works for UNet but gives CUDA out of memory error with DynUNet and SwinUNETR

[ahxmeds]

Hi,

I have used the 3D segmentation code (with modification) based on the spleen segmentation code in the MONAI tutorials for my own PET/CT segmentation project. My code seems to work fine when I use UNet as the network. But if I try using others network (like SwinUNETR or DynUNet), I get RuntimeError: CUDA out of memory. I am only testing this with 2 training and validation PET/CT images for now.

My train_transform and valid_transform look like this:

spatialsize=(192, 192, 192)
mod_keys = ['CT', 'PT', 'GT']
train_transforms = Compose(
    [
        LoadImaged(keys=mod_keys),
        EnsureChannelFirstd(keys=mod_keys),
        ScaleIntensityRanged(keys=['CT'], a_min=-1024, a_max=1024, b_min=0, b_max=1, clip=True),
        ScaleIntensityd(keys=['PT'], minv=0, maxv=1),
        CropForegroundd(keys=mod_keys, source_key='PT'),
        Orientationd(keys=mod_keys, axcodes="RAS"),
        Spacingd(keys=mod_keys, pixdim=(2.0, 2.0, 2.0), mode=('bilinear', 'bilinear', 'nearest')),
        RandCropByPosNegLabeld(
            keys=mod_keys,
            label_key='GT',
            spatial_size = spatialsize,
            pos=4,
            neg=1,
            num_samples=1,
            image_key='PT',
            image_threshold=0,
        ),
        RandAffined(
            keys=mod_keys,
            mode=('bilinear', 'bilinear', 'nearest'),
            prob=0.5,
            spatial_size = spatialsize,
            translate_range=(10,10,10),
            rotate_range=(0, 0, np.pi/15),
            scale_range=(0.1, 0.1, 0.1)),
        Rand3DElasticd(
            keys=mod_keys,
            sigma_range=(0.0, 1.0),
            magnitude_range=(0.0, 1.0),
            spatial_size = spatialsize,
            prob=0.4,
        ),
        ConcatItemsd(keys=['CT', 'PT'], name='CTPT', dim=0),
        DeleteItemsd(keys=['CT', 'PT'])
    ]
)
# %%
valid_transforms = Compose(
    [
        LoadImaged(keys=mod_keys),
        EnsureChannelFirstd(keys=mod_keys),
        ScaleIntensityRanged(keys=['CT'], a_min=-1024, a_max=1024, b_min=0, b_max=1, clip=True),
        ScaleIntensityd(keys=['PT'], minv=0, maxv=1),
        CropForegroundd(keys=mod_keys, source_key='PT'),
        Orientationd(keys=mod_keys, axcodes="RAS"),
        Spacingd(keys=mod_keys, pixdim=(2.0, 2.0, 2.0), mode=('bilinear', 'bilinear', 'nearest')),
        ConcatItemsd(keys=['CT', 'PT'], name='CTPT', dim=0),
        DeleteItemsd(keys=['CT', 'PT'])
    ]
)
#%%
dataset_train = CacheDataset(data=train_data[0:2], transform=train_transforms, cache_rate=1, num_workers=16)
dataloader_train = DataLoader(dataset_train, batch_size=1, shuffle=True, num_workers=16)

#%%
dataset_valid = CacheDataset(data=valid_data[0:2], transform=valid_transforms, cache_rate=1, num_workers=16)
dataloader_valid = DataLoader(dataset_valid, batch_size=1, shuffle=False, num_workers=16)

And I am doing the model definition like this:

device = torch.device('cuda:0')
# model = UNet(   
#     spatial_dims=3,
#     in_channels=2,
#     out_channels=2,
#     channels=(16, 32, 64, 128, 256, 512),
#     strides=(2,2,2,2,2),
#     num_res_units=2
# ).to(device)

model = SwinUNETR(
    img_size=spatialsize,
    in_channels=2,
    out_channels=2,
    feature_size=48,
    drop_rate=0.0,
    attn_drop_rate=0.0,
    dropout_path_rate=0.0,
    use_checkpoint=False,
).to(device)
model = nn.DataParallel(model, device_ids=[0, 1, 2, 3])

loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=False, reduction="mean")
scheduler = CosineAnnealingLR(optimizer, T_max=1200, eta_min=0)

The training for the UNet model (that has been commented out above) seems to work with this code, but using SwinUNETR (or DynUNet) generates this error:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [88], in ()
     13 inputs, labels = (
     14     batch_data['CTPT'].to(device),
     15     batch_data['GT'].to(device),
     16 )
     17 optimizer.zero_grad()
---> 18 outputs = model(inputs)
     19 loss = loss_function(outputs, labels)
     20 loss.backward()

File /anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py:1110, in Module._call_impl(self, *input, **kwargs)
   [1106](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1105) # If we don't have any hooks, we want to skip the rest of the logic in
   [1107](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1106) # this function, and just call forward.
   [1108](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1107) if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   [1109](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1108)         or _global_forward_hooks or _global_forward_pre_hooks):
-> [1110](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1109)     return forward_call(*input, **kwargs)
   [1111](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1110) # Do not call functions when jit is used
   [1112](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/modules/module.py?line=1111) full_backward_hooks, non_full_backward_hooks = [], []

File /anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py:168, in DataParallel.forward(self, *inputs, **kwargs)
    [166](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py?line=165)     return self.module(*inputs[0], **kwargs[0])
    [167](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py?line=166) replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
--> [168](file:///anaconda/envs/dlml/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py?line=167) outputs = self.parallel_apply(replicas, inputs, kwargs)
...
    attn_mask = compute_mask([dp, hp, wp], window_size, shift_size, x.device)
  File "/anaconda/envs/dlml/lib/python3.8/site-packages/monai/networks/nets/swin_unetr.py", line 793, in compute_mask
    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
RuntimeError: CUDA out of memory. Tried to allocate 1.20 GiB (GPU 0; 15.78 GiB total capacity; 13.15 GiB already allocated; 280.69 MiB free; 14.30 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

Surprisingly, I tried running the whole thing on the CPU and the training runs successfully on the CPU with absolutely no errors (although very slow, taking about 2 mins per epoch, even with 2 training images). Please let me know how to fix this error (as I do not understand it) and help me run my code on the GPU(s).

[KumoLiu]

Hi @ahxmeds, since SwinUNETR will use more memory than UNet, you can run with UNet but not SwinUNETR is reasonable. I noticed that you are using RandCropByPosNegLabeld in the preprocessing, I would suggest you crop smaller patches such as (96, 96, 96) if the size makes sense. Hope it can help you, thanks!

[ahxmeds]

@KumoLiu Thank you. That was helpful, the SwinUNETR training code does run with training patch size (96, 96, 96) and inference window size of (96, 96, 96) (for validation during training). But, my question still remains. Let me give you a little bit of context regarding my experiments. I was basically training different UNet models using different patch sizes of size (N, N, N), where N = {96, 128, 160, 192, 224, 256}. For each of these training patches, I was using an inference sliding window size of (M, M, M), where M = {96, 128, 160, 192, 224, 256, 288}. Basically, the (N, M) pairs are my hyperparameters, which I am optimizing on my validation set during training. Here are some plots for training patch sizes N=96 and N=192 for all inference window sizes (M) using UNet:

As you can see in the above plots, for N=96, the optimal M=256, and for N=192, M=192. So, it is evident that the training patch size and the corresponding inference window can strongly influence the model performance and model selection. And, my initial plan was to replicate this ablation experiments in other networks as well, like SwinUNETR or 'DynUNet'.

I also noticed another thing: I used from torchsummary import summary and the number of trainable parameters for UNet is independent of the spatial_size, whereas for DynUNet and 'SwinUNETR`, the number of parameters increase with increasing training patch size, which is why I cannot accommodate training these in my 4 GPUs (16 GiB RAM each).

I tried training SwinUNETR with (N, M)=(96,96) and (although it is still training) it seems to perform only as good as UNet with (N, M)=(96,96), whereas, I know that UNet with (N, M)=(192,192) has much better performance (see training and validation curves below). So in that regard, is there any merit to using networks other than UNet for my project given my hardware constraints at this moment?

Training loss and validation DSC for UNet with (N, M)=(96,96):

Training loss and validation DSC for UNet with (N, M)=(192,192):

Training loss and validation DSC for SwinUNETR with (N, M)=(96,96):

Again, thanks for your help so far. :)