spleen segmentation 3d, RuntimeError: CUDA error: device-side assert triggered

[ctpn]

Hello,

I am trying to modify the 3D Spleen Segmentation tutorial for my own dataset. In the 'Execute a typical PyTorch training process' cells, I get the error, RuntimeError: CUDA error: device-side assert triggered.

From StackOverflow, it seems that this could is likely a mismatch between the number of channels in 'outputs' vs 'labels'. My labelled images are grayscale 3D images, and I have kept out_channels=2 when defining the UNet model. Not sure if this is relevant, but I have modified the patch size, channels, and changed the z-axis spacing (Spacingd) based on this issue, because my images are shallow (e.g. 512x512x14).

image shape: torch.Size([512, 512, 34]), label shape: torch.Size([512, 512, 34])
image dim: 3, label dim: 3

Other users have also recommended adding os.environ['CUDA_LAUNCH_BLOCKING'] = '1' to get a more detailed error message, however this has not changed the error messages for me. Another suggestion was to run on CPU, but this colab notebook seems to only run on GPU.

Any advice would be greatly appreciated, thank you. Code is also included below.

Error message:

----------
epoch 1/20
torch.Size([4, 1, 16, 16, 16])
batch_data image: torch.Size([4, 1, 16, 16, 16])
batch_data label: torch.Size([4, 1, 16, 16, 16])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-7f47132bac00> in <module>()
     27         print(f"batch_data label: {batch_data['label'].shape}")
     28         outputs = model(inputs)
---> 29         loss = loss_function(outputs, labels)
     30         loss.backward() # calculate gradient
     31         optimizer.step() # update weights

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.6/dist-packages/monai/losses/dice.py in forward(self, input, target)
    145             reduce_axis = [0] + reduce_axis
    146 
--> 147         intersection = torch.sum(target * input, dim=reduce_axis)
    148 
    149         if self.squared_pred:

RuntimeError: CUDA error: device-side assert triggered

Code:

%pip install -q "monai[gdown, nibabel, tqdm, ignite]"
%pip install -q "monai[itk]"
%pip install git+https://github.com/Project-MONAI/MONAI#egg=MONAI
%pip install -q matplotlib
%matplotlib inline

import glob
import os
import shutil
import tempfile
import itk
import matplotlib.pyplot as plt
import torch

from monai.apps import download_and_extract
from monai.config import print_config
from monai.data import CacheDataset, DataLoader, Dataset
from monai.inferers import sliding_window_inference
from monai.losses import DiceLoss
from monai.metrics import compute_meandice
from monai.networks.layers import Norm
from monai.networks.nets import UNet
from monai.transforms import (
    AsDiscrete,
    AddChanneld,
    Compose,
    CropForegroundd,
    LoadImaged,
    Orientationd,
    RandCropByPosNegLabeld,
    ScaleIntensityRanged,
    Spacingd,
    ToTensord,
)
from monai.utils import first, set_determinism
print_config()

root_dir2 = '/content'
data_dir2 = root_dir2
train_images2 = sorted(glob.glob(os.path.join(data_dir2, "*x.nii.gz")))
train_labels2 = sorted(glob.glob(os.path.join(data_dir2, "*y.nii.gz")))
data_dicts2 = [
    {"image": image_name, "label": label_name}
    for image_name, label_name in zip(train_images2, train_labels2)
]
train_files2, val_files2 = data_dicts2[:-2], data_dicts2[-2:] 

set_determinism(seed=0)

train_transforms2 = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(1.0, 1.0, 0.24), mode=("bilinear", "nearest")),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        ScaleIntensityRanged(
            keys=["image"], a_min=0, a_max=255, b_min=0.0, b_max=1.0, clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(16, 16, 16),
            pos=1,
            neg=1,
            num_samples=4,
            image_key="image",
            image_threshold=0,
        ),
        ToTensord(keys=["image", "label"]),
        # user can also add other random transforms
        # RandAffined(keys=['image', 'label'], mode=('bilinear', 'nearest'), prob=1.0, spatial_size=(96, 96, 96),
        #             rotate_range=(0, 0, np.pi/15), scale_range=(0.1, 0.1, 0.1)),
        
    ]
)
val_transforms2 = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        AddChanneld(keys=["image", "label"]),
        Spacingd(keys=["image", "label"], pixdim=(1, 1, 0.24), mode=("bilinear", "nearest")),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        ScaleIntensityRanged(
            keys=["image"], a_min=0, a_max=255, b_min=0.0, b_max=1.0, clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        ToTensord(keys=["image", "label"]),
    ]
)

check_ds = Dataset(data=val_files2, transform=val_transforms2)
check_loader = DataLoader(check_ds, batch_size=1)
check_data = first(check_loader)
#check_data = check_loader[0]
image, label = (check_data["image"][0][0], check_data["label"][0][0])
print(f"image shape: {image.shape}, label shape: {label.shape}")
print(f"image dim: {image.ndim}, label dim: {label.ndim}")
# plot the slice [:, :, 80]
plt.figure("check", (12, 6))
plt.subplot(1, 2, 1)
plt.title("image")
plt.imshow(image[:, :, 1], cmap="gray")
plt.subplot(1, 2, 2)
plt.title("label")
plt.imshow(label[:, :, 1], cmap='gray')
plt.show()

train_ds2 = CacheDataset(data=train_files2, transform=train_transforms2, cache_rate=1.0, num_workers=4)
train_loader2 = DataLoader(train_ds2, batch_size=1, shuffle=True, num_workers=4)
val_ds2 = CacheDataset(data=val_files2, transform=val_transforms2, cache_rate=1.0, num_workers=4)
val_loader2 = DataLoader(val_ds2, batch_size=1, num_workers=4)

device = torch.device("cuda:0")
model = UNet(
    dimensions=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 128, 256),
    strides=(2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True) # the higher the better
optimizer = torch.optim.Adam(model.parameters(), 1e-4)

epoch_num = 20
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
post_pred = AsDiscrete(argmax=True, to_onehot=True, n_classes=2)
post_label = AsDiscrete(to_onehot=True, n_classes=2)

for epoch in range(epoch_num):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{epoch_num}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader2:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        print(inputs.shape)
        optimizer.zero_grad()
        print(f"batch_data image: {batch_data['image'].shape}")
        print(f"batch_data label: {batch_data['label'].shape}")
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward() # calculate gradient
        optimizer.step() # update weights
        epoch_loss += loss.item() # sum epoch losses
        print(f"{step}/{len(train_ds2) // train_loader2.batch_size}, train_loss: {loss.item():.4f}")
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            metric_sum = 0.0
            metric_count = 0
            for val_data2 in val_loader2:
                val_inputs2, val_labels2 = (
                    val_data2["image"].to(device),
                    val_data2["label"].to(device),
                )
                roi_size = (128, 128, 128)
                sw_batch_size = 4
                val_outputs2 = sliding_window_inference(val_inputs2, roi_size, sw_batch_size, model)
                val_outputs2 = post_pred(val_outputs2)
                val_labels2 = post_label(val_labels2)
                value = compute_meandice(
                    y_pred2=val_outputs2,
                    y2=val_labels2,
                    include_background=False,
                )
                metric_count += len(value)
                metric_sum += value.sum().item()
            metric = metric_sum / metric_count
            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(root_dir, "best_metric_model.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f"\nbest mean dice: {best_metric:.4f} at epoch: {best_metric_epoch}"
            )

[rijobro]

Hi, before looking through your code more closely, if the StackOverflow post recommends running on the CPU rather than the GPU, then this can be done by replacing this: device = torch.device("cuda:0") with this: device = torch.device("cpu").

Perhaps try that first and see if the error is more informative.

[rijobro]

Could you put the following after outputs = model(inputs):

print(f"output shape: {output.shape}")

And also print all the unique values in your label image?

torch.unique(label)

[ctpn]

Thank you very much for your reply. After running it on CPU, the error message was indeed more informative:

RuntimeError                              Traceback (most recent call last)
<ipython-input-18-5a9c1209c649> in <module>()
     30         print(f"label shape: {label.shape}")
     31         print(f"label image unique values: {torch.unique(label)}")
---> 32         loss = loss_function(outputs, labels)
     33         loss.backward() # calculate gradient
     34         optimizer.step() # update weights

/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    725             result = self._slow_forward(*input, **kwargs)
    726         else:
--> 727             result = self.forward(*input, **kwargs)
    728         for hook in itertools.chain(
    729                 _global_forward_hooks.values(),

/usr/local/lib/python3.6/dist-packages/monai/losses/dice.py in forward(self, input, target)
    125                 warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
    126             else:
--> 127                 target = one_hot(target, num_classes=n_pred_ch)
    128 
    129         if not self.include_background:

/usr/local/lib/python3.6/dist-packages/monai/networks/utils.py in one_hot(labels, num_classes, dtype, dim)
     56 
     57     o = torch.zeros(size=sh, dtype=dtype, device=labels.device)
---> 58     labels = o.scatter_(dim=dim, index=labels.long(), value=1)
     59 
     60     return labels

RuntimeError: index 255 is out of bounds for dimension 1 with size 2

Adding the 2 lines of code you recommended:

torch.Size([4, 1, 16, 16, 16]) 
batch_data image: torch.Size([4, 1, 16, 16, 16]) 
batch_data label: torch.Size([4, 1, 16, 16, 16])
output shape: torch.Size([4, 2, 16, 16, 16])
label shape: torch.Size([512, 512, 34])
label image unique values: tensor([  0., 255.])

I then converted my labelled image from [0,255] to [0,1], and it was resolved. Output:

torch.Size([4, 1, 16, 16, 16])
batch_data image: torch.Size([4, 1, 16, 16, 16])
batch_data label: torch.Size([4, 1, 16, 16, 16])
outputs shape: torch.Size([4, 2, 16, 16, 16])
label shape: torch.Size([512, 512, 34])
label image unique values: tensor([0., 1.])

Much appreciated, thank you again!