Catch Invertd Transform if CUDA Out of Memory

[idinsmore1]

Hello, thanks for such an amazing job with MONAI.
I'm running an inference pipeline where I want to do the postprocessing on the GPU when possible, but catch and switch to CPU if out of memory. My preprocessing and postprocessing transformations always work when not OOM, but when the exception catches, I always get the same error.

This is my code

preprocessing = Compose([
            LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            ThresholdIntensityd(keys=['image'], threshold=task_config['percentile_95'], above=False, cval=task_config['percentile_95']),
            ThresholdIntensityd(keys=['image'], threshold=task_config['percentile_05'], above=True, cval=task_config['percentile_05']),
            NormalizeIntensityd(keys=['image'], subtrahend=task_config['mean'], divisor=task_config['std']),
            CropForegroundd(keys=["image"], source_key="image", allow_smaller=True, select_fn=lambda x: x > task_config['crop_threshold']),
            Orientationd(keys=['image'], axcodes='RAS'),
            Spacingd(keys=['image'], pixdim=task_config['spacing'], mode='bilinear'),
            EnsureTyped(keys=['image'], track_meta=True)
        ])
postprocessing_transform = Compose([
      # EnsureTyped(keys=['pred'], device=postprocessing_device),
      Activationsd(keys=['pred'], softmax=True),
      AsDiscreted(keys=['pred'], argmax=True),
      Invertd(keys=['pred'], transform=preprocessing, orig_keys='image', meta_keys='image_meta_dict', nearest_interp=True, to_tensor=True),
      SqueezeDimd(keys=['pred'], dim=0),
      ToNumpyd(keys=['pred'], dtype=np.uint8)
  ])
gpu_postprocessing = Compose([EnsureTyped(keys=['pred'], device=device), postprocessing_transform])
cpu_postprocessing = Compose([EnsureTyped(keys=['pred'], device='cpu'), postprocessing_transform])
#     logger.info(f'Changing batch size to {batch_size} for larger patch size')
inferer = SlidingWindowInfererAdapt(roi_size=model_config[task]['patch_size'], sw_batch_size=batch_size, overlap=0.5)
# Run the inference loop
with autocast():
    with torch.no_grad():
        for data in dataloader:
            images = data['image'].to(device)
            # Run inference
            start_time = time.time()
            pred = inferer(inputs=images, network=model)
            inference_time = round(time.time() - start_time, 2)
            data['pred'] = pred
            # Delete the images to save gpu memory
            del images 
            # Run postprocessing. Only have one item so take index 0
            processing_start = time.time()
            # Attempt to run postprocessing on GPU, if it fails due to OOM, run it on CPU
            # If the prediction is on CPU (== -1), we just go right to CPU postprocessing
            if data['pred'].get_device() != -1:
                try:
                    out = [gpu_postprocessing(i) for i in decollate_batch(data)][0]
                except RuntimeError as e:
                    # print(e)
                    # print(f'{data=}')
                    print(f'GPU Postprocessing failed due to CudaOOM for {data["metadata"]["original_name"]}. Switching to CPU')
                    data['pred'] = pred.cpu()
                    out = [cpu_postprocessing(i) for i in decollate_batch(data)][0]
            else:
                out = [cpu_postprocessing(i) for i in decollate_batch(data)][0]

[KumoLiu]

Hi @idinsmore1, could you please check the env var MONAIEnvVars.trace_transform()? It should be 1 if you want to trace transforms. BTW, by default it's 1.

Thanks.

[idinsmore1]

Hi @KumoLiu I have checked this, it's still 1 after the error catches, and looking at data['pred'].applied_operations and data['image'].applied_operations, all the transformations have tracing = True, for some reason after the error it's no longer recognized and I'm not quite sure why. I also tried rerunning the inference inside the error exception catch, the same error occurs when trying to run the postprocessing operations.

[idinsmore1]

Hi there, I'm going to reopen this discussion, as this operation does not seem to be consistently working and I'm more confused than before. I think there may be something going on with the storage/collection of the metadata when the gpu_postprocessing fails, and the cpu_postprocessing tries to pick the process back up. If anyone has any insight to this, it would be greatly appreciated. thanks!

[idinsmore1]

building on top of this, something is definitely going on with metadata after attempting to just skip postprocessing inside the error handling altogether. When I try the following code, everything works before the first instance of the error, but after that, every following image in the dataloader fails due to the same tracing error.

preprocessing = Compose([
            LoadImaged(keys=['image']),
            EnsureChannelFirstd(keys=['image']),
            ThresholdIntensityd(keys=['image'], threshold=task_config['percentile_95'], above=False, cval=task_config['percentile_95']),
            ThresholdIntensityd(keys=['image'], threshold=task_config['percentile_05'], above=True, cval=task_config['percentile_05']),
            NormalizeIntensityd(keys=['image'], subtrahend=task_config['mean'], divisor=task_config['std']),
            CropForegroundd(keys=["image"], source_key="image", allow_smaller=True, select_fn=lambda x: x > task_config['crop_threshold']),
            Orientationd(keys=['image'], axcodes='RAS'),
            Spacingd(keys=['image'], pixdim=task_config['spacing'], mode='bilinear'),
            EnsureTyped(keys=['image'], track_meta=True)
        ])
postprocessing_transform = Compose([
      # EnsureTyped(keys=['pred'], device=postprocessing_device),
      Activationsd(keys=['pred'], softmax=True),
      AsDiscreted(keys=['pred'], argmax=True),
      Invertd(keys=['pred'], transform=preprocessing, orig_keys='image', meta_keys='image_meta_dict', nearest_interp=True, to_tensor=True),
      SqueezeDimd(keys=['pred'], dim=0),
      ToNumpyd(keys=['pred'], dtype=np.uint8)
  ])
gpu_postprocessing = Compose([EnsureTyped(keys=['pred'], device=device), postprocessing_transform])
cpu_postprocessing = Compose([EnsureTyped(keys=['pred'], device='cpu'), postprocessing_transform])
#     logger.info(f'Changing batch size to {batch_size} for larger patch size')
inferer = SlidingWindowInfererAdapt(roi_size=model_config[task]['patch_size'], sw_batch_size=batch_size, overlap=0.5)
# Run the inference loop
with autocast():
    with torch.no_grad():
        for data in dataloader:
            images = data['image'].to(device)
            # Run inference
            start_time = time.time()
            pred = inferer(inputs=images, network=model)
            inference_time = round(time.time() - start_time, 2)
            data['pred'] = pred
            # Delete the images to save gpu memory
            del images 
            # Run postprocessing. Only have one item so take index 0
            processing_start = time.time()
            # Attempt to run postprocessing on GPU, if it fails due to OOM, run it on CPU
            # If the prediction is on CPU (== -1), we just go right to CPU postprocessing
            if data['pred'].get_device() != -1:
                try:
                    out = [gpu_postprocessing(i) for i in decollate_batch(data)][0]
                except RuntimeError as e:
                    print('Postprocessing failed') # After the first occurence of this error catch, everything breaks
            else:
                out = [cpu_postprocessing(i) for i in decollate_batch(data)][0]