Error while training on the BraTS 2023 dataset

[amida47]

the following is code I used to train on the Brats 2023 dataset

import os
import shutil
import tempfile
import time
import matplotlib.pyplot as plt
from monai.apps import DecathlonDataset
from monai.config import print_config
from monai.data import DataLoader, decollate_batch
from monai.handlers.utils import from_engine
from monai.losses import DiceLoss
from monai.inferers import sliding_window_inference
from monai.metrics import DiceMetric
from monai.networks.nets import SegResNet
from monai.transforms import (
    Activations,
    Activationsd,
    AsDiscrete,
    AsDiscreted,
    Compose,
    Invertd,
    LoadImaged,
    MapTransform,
    NormalizeIntensityd,
    Orientationd,
    RandFlipd,
    RandScaleIntensityd,
    RandShiftIntensityd,
    RandSpatialCropd,
    Spacingd,
    EnsureTyped,
    EnsureChannelFirstd,
)
from monai.utils import set_determinism

import torch
class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
    """
    Convert labels to multi channels based on brats classes:
    label 1 is the peritumoral edema
    label 2 is the GD-enhancing tumor
    label 3 is the necrotic and non-enhancing tumor core
    The possible classes are TC (Tumor core), WT (Whole tumor)
    and ET (Enhancing tumor).

    """

    def __call__(self, data):
        d = dict(data)
        for key in self.keys:
            result = []
            # merge label 2 and label 3 to construct TC
            result.append(torch.logical_or(d[key] == 2, d[key] == 3))
            # merge labels 1, 2 and 3 to construct WT
            result.append(torch.logical_or(torch.logical_or(d[key] == 2, d[key] == 3), d[key] == 1))
            # label 2 is ET
            result.append(d[key] == 2)
            d[key] = torch.stack(result, axis=0).float()
            print(f"after classes creation {d[key].shape}")
        return d

# this is a custom loading transform I had to write because I couldn't find one that can load 4 images and combine them into one
from monai.data import ITKReader
import time
import numpy as np

reader = ITKReader()

class Load4Modalitiesd(MapTransform):
    def __call__(self, data):
        d = dict(data)
        # create empty dictionary
        out = {}
        # calculate and print the time it takes to load the data in seconds
        
        start_time = time.time()
        for key in self.keys:
            if len(d[key]) == 4:
                print(d[key])
                out['image'] = reader.get_data(reader.read(d[key]))[0]
                out['image'] = np.transpose(out['image'], (0, 3, 1, 2))
                print(f"image shape: {out['image'].shape}")
            else:

                out['label'] = reader.get_data(reader.read(d[key]))[0]
                print(f"seg shape: {out['label'].shape}")

                out['label'] = np.transpose(out['label'], (2, 0, 1))
        print(f"Loading image took {time.time() - start_time} seconds")

        return out
train_transform = Compose(
    [
        # load 4 Nifti images and stack them together
        Load4Modalitiesd(keys=["imgs", "seg"]),
        # EnsureChannelFirstd(keys="image"),
        EnsureTyped(keys=["image", "label"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.0, 1.0, 1.0),
            mode=("bilinear", "nearest"),
        ),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
        RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
    ]
)
val_transform = Compose(
    [
        Load4Modalitiesd(keys=["imgs", "seg"]),
        # EnsureChannelFirstd(keys="image"),
        EnsureTyped(keys=["image", "label"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.0, 1.0, 1.0),
            mode=("bilinear", "nearest"),
        ),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
    ]
)
# creating the list of the files
from glob import glob

tempdir = '/content/ASNR-MICCAI-BraTS2023-GLI-Challenge-TrainingData'


images_t1c = sorted(glob(os.path.join(tempdir, "*/*-t1c.nii.gz")))
images_t2f = sorted(glob(os.path.join(tempdir, "*/*-t2f.nii.gz")))
images_t1n = sorted(glob(os.path.join(tempdir, "*/*-t1n.nii.gz")))
images_t2w = sorted(glob(os.path.join(tempdir, "*/*-t2w.nii.gz")))
print(f"image files: {len(images_t1c)}")

segs = sorted(glob(os.path.join(tempdir, "*/*seg.nii.gz")))

train_files = [{"imgs": [img,img2,img3,img4], "seg": seg} for img,img2,img3,img4, seg in zip(images_t2f[:20],
                                                            images_t1c[:20],
                                                            images_t1n[:20],
                                                            images_t2w[:20],
                                                            segs[:20])]
val_files = [{"imgs": [img,img2,img3,img4], "seg": seg} for img,img2,img3,img4, seg in zip(images_t2f[20:], 
                                                            images_t1c[20:],
                                                            images_t1n[20:],
                                                            images_t2w[20:],
                                                            segs[20:])]

from PIL import Image
from torch.utils.tensorboard import SummaryWriter
from monai.visualize import plot_2d_or_3d_image

# create a training data loader
train_ds = monai.data.Dataset(data=train_files, transform=train_transform)
# use batch_size=2 to load images and use RandCropByPosNegLabeld to generate 2 x 4 images for network training
train_loader = DataLoader(
    train_ds,
    batch_size=2,
    shuffle=True,
    num_workers=0,
    collate_fn=list_data_collate,
    pin_memory=torch.cuda.is_available(),
)
# create a validation data loader
val_ds = monai.data.Dataset(data=val_files, transform=val_transform)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=0, collate_fn=list_data_collate)
dice_metric = DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
post_trans = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
# create UNet, DiceLoss and Adam optimizer
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = monai.networks.nets.UNet(
    spatial_dims=3,
    in_channels=4,
    out_channels=3,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)
loss_function = monai.losses.DiceLoss(sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-3)

# start a typical PyTorch training
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = list()
metric_values = list()
writer = SummaryWriter()
for epoch in range(10):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{10}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        epoch_len = len(train_ds) // train_loader.batch_size
        print(f"{step}/{epoch_len}, train_loss: {loss.item():.4f}")
        writer.add_scalar("train_loss", loss.item(), epoch_len * epoch + step)
    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():
            val_images = None
            val_labels = None
            val_outputs = None
            for val_data in val_loader:
                val_images, val_labels = val_data["image"].to(device), val_data["label"].to(device)
                roi_size = (96, 96)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(val_images, roi_size, sw_batch_size, model)
                val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                # compute metric for current iteration
                dice_metric(y_pred=val_outputs, y=val_labels)
            # aggregate the final mean dice result
            metric = dice_metric.aggregate().item()
            # reset the status for next validation round
            dice_metric.reset()
            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), "best_metric_model_segmentation2d_dict.pth")
                print("saved new best metric model")
            print(
                "current epoch: {} current mean dice: {:.4f} best mean dice: {:.4f} at epoch {}".format(
                    epoch + 1, metric, best_metric, best_metric_epoch
                )
            )
            writer.add_scalar("val_mean_dice", metric, epoch + 1)
            # plot the last model output as GIF image in TensorBoard with the corresponding image and label
            plot_2d_or_3d_image(val_images, epoch + 1, writer, index=0, tag="image")
            plot_2d_or_3d_image(val_labels, epoch + 1, writer, index=0, tag="label")
            plot_2d_or_3d_image(val_outputs, epoch + 1, writer, index=0, tag="output")

print(f"train completed, best_metric: {best_metric:.4f} at epoch: {best_metric_epoch}")
writer.close()

so to summarize, I fetch all the files in the dataset folder, the data is a list of dictionaries in each dict (corresponds to a patient) there is two keys "imgs": a list of the 4 modalities , and "seg" is the segmentation map of that patient .
I load the 4 modalities and stack them to get a shape of (4, 240, 240, 155) then permute to (4, 155, 240, 240), I do the same to the segmentation map load and permute to get a shape of (155, 240, 240), this is the only thing I added by myself
when the training loop starts I get this error

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
[<ipython-input-15-d7a9b4140ddb>](https://localhost:8080/#) in <cell line: 41>()
     49         inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
     50         optimizer.zero_grad()
---> 51         outputs = model(inputs)
     52         loss = loss_function(outputs, labels)
     53         loss.backward()

11 frames
[/usr/local/lib/python3.10/dist-packages/torch/_tensor.py](https://localhost:8080/#) in __torch_function__(cls, func, types, args, kwargs)
   1293 
   1294         with _C.DisableTorchFunctionSubclass():
-> 1295             ret = func(*args, **kwargs)
   1296             if func in get_default_nowrap_functions():
   1297                 return ret

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 39 but got size 40 for tensor number 1 in the list.

for the full traceback visit https://pastebin.com/zZSRihsQ

[KumoLiu]

Hi @amida47, could you please check the shape of your data, looks like you have a mismatched shape? Thanks!

[KumoLiu]

Yes, thanks for clarifying. Let's check the shape when a crash happens, like before outputs = model(inputs)?

[amida47]

....
for batch_data in train_loader:
        step += 1
        inputs, labels = batch_data["image"].to(device), batch_data["label"].to(device)
        optimizer.zero_grad()
        print(f"inputs shape ====> {inputs.shape}")
        print(f"labels shape ====> {labels.shape}")
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
....

these are shapes before the model call

inputs shape ====> torch.Size([2, 4, 155, 240, 240])
labels shape ====> torch.Size([2, 3, 155, 240, 240])

somehow I think the image resolution being 240x240 is the cause, I will try to pad that to 256.
EDIT: tried the pad thing and didn't work

[KumoLiu]

ok, perhaps you should ensure the shape of your data is a multiple of 16 (including 155 and 240).

[amida47]

well 240 is easy we just pad the image, but 155 is number of mri slices, do you suggest to just add empty slices?

[KumoLiu]

Yes, you could use ResizeWithPadOrCrop or if you don't want to pad the image, you can also use Resize with the appropriate mode for the image and label.

MONAI/monai/transforms/croppad/array.py

Line 1405 in 0a3cdd9

class ResizeWithPadOrCrop(InvertibleTransform, LazyTransform):

MONAI/monai/transforms/spatial/array.py

Line 725 in 0a3cdd9

mode: {``"nearest"``, ``"nearest-exact"``, ``"linear"``, ``"bilinear"``, ``"bicubic"``, ``"trilinear"``, ``"area"``}

[wyli]

the original example has the cropping transform RandSpatialCropd(keys=["image", "label"], roi_size=[224, 224, 144], random_size=False),

https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/brats_segmentation_3d.ipynb

batch_size=2 means stacking two input volumes at the batch dim, without the transform, the volumes will have different spatial sizes and could not be stacked together. so you can either add the crop/pad transform to ensure consistent spatial size or change the batch size to 1.

[amida47]

thanks for the response.
just wanted to point out that I made sure to delete it from my code, as I didn't want the random crop, as you can see in the code above, I just forgot to delete RandSpatialCropd from the imports.

train_transform = Compose(
    [
        # load 4 Nifti images and stack them together
        Load4Modalitiesd(keys=["imgs", "seg"]),
        # EnsureChannelFirstd(keys="image"),
        EnsureTyped(keys=["image", "label"]),
        ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.0, 1.0, 1.0),
            mode=("bilinear", "nearest"),
        ),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
        RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
        NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
        RandScaleIntensityd(keys="image", factors=0.1, prob=1.0),
        RandShiftIntensityd(keys="image", offsets=0.1, prob=1.0),
    ]
)