Error: received 0 items of ancdata

[jmenville11]

I'm building a U-Net image segmentation model based off of MONAI's BRATS tutorial. I've gotten the model to run successfully and have started playing with the model parameters to see how I could optimize training speed. I've found that whenever I reduce Cache Rate to anything less than 1.0, the model will go through 4-5 epochs before failing. I get "Runtime Error: received 0 items of ancdata".

I've read of others having similar issues using PyTorch, but haven't been able to identify the cause. Should this variable not be changed?

Code is below:

Setup imports

    import os
    import shutil
    import tempfile
    import glob
    import nibabel as nib
    import torch
    from pathlib import Path
    import matplotlib.pyplot as plt
    import numpy as np
    from numpy import array
    from monai.apps import DecathlonDataset
    from monai.data import CacheDataset, DataLoader, decollate_batch, Dataset, NiftiSaver
    from monai.config import print_config
    from monai.data import DataLoader
    from monai.losses import DiceLoss
    from monai.metrics import DiceMetric
    from monai.networks.nets import UNet
    from monai.transforms import (
        Activations,
        AsChannelFirstd,
        AsDiscrete,
        CenterSpatialCropd,
        Compose,
        LoadImage,
        LoadImaged,
        MapTransform,
        NormalizeIntensityd,
        Orientationd,
        RandAffined,
        RandFlipd,
        RandScaleIntensityd,
        RandShiftIntensityd,
        RandSpatialCropd,
        ResizeWithPadOrCropd,
        Spacingd,
        SpatialPadd,
        ToTensord,
        ToTensor,    
        RandCropByPosNegLabeld,
        AddChanneld,
    )
    from monai.utils import set_determinism
    print_config()

Setup data directory

    data_dir = Path("/gpfs/home/jmenvill/ICH/Model_Directory_Uniform_Sizing_2")
    model_dir = Path("/users/jmenvill/ICH/Models")
    os.chdir(data_dir)

Establish file paths

    train_images = sorted(glob.glob(os.path.join(data_dir, "train", "Images", "*_FLAIR.nii")))
    train_labels = sorted(glob.glob(os.path.join(data_dir, "train", "Labels", "*_label.nii")))
    val_images = sorted(glob.glob(os.path.join(data_dir, "val", "Images", "*_FLAIR.nii")))
    val_labels = sorted(glob.glob(os.path.join(data_dir, "val", "Labels", "*_label.nii")))
    test_images = sorted(glob.glob(os.path.join(data_dir, "test", "Images", "*_FLAIR.nii")))
    
    print('Training images found: ',len(train_images))
    print('Training segmentations found: ',len(train_labels))
    print('Validation images found: ',len(val_images))
    print('Validation segmentations found: ',len(val_labels))

Zip labels and images together into a dictionary for both training and validation sets

    train_files = [
        {"image": image_name, "label": label_name}
        for image_name, label_name in zip(train_images, train_labels)
    ]
    val_files = [
        {"image": image_name, "label": label_name}
        for image_name, label_name in zip(val_images, val_labels)
    ]

Set deterministic training for reproducibility

    set_determinism(seed=0)

Convert brain tumor labels

    class ConvertToMultiChannelBasedOnBratsClassesd(MapTransform):
        def __call__(self, data):
            d = dict(data)
            for key in self.keys:
                result = []
                # label 1 is HE (hematoma)
                result.append(d[key] == 1)
                # label 2 is ED (edema)
                result.append(d[key] == 2)
                d[key] = np.stack(result, axis=0).astype(np.float32)
            return d  

Setup transforms for training and validation

    train_transform = Compose(
        [
            LoadImaged(keys=["image", "label"]),
            AddChanneld(keys=["image"]),
            ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
            Spacingd(
                keys=["image", "label"],
                pixdim=(1,1,1),
                mode=("bilinear", "nearest"),
            ),
            Orientationd(
                keys=["image", "label"],
                axcodes="LPS"
            ), 
            ResizeWithPadOrCropd(keys=["image", "label"],
                spatial_size=(256,256,256),
                mode="constant"),
            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            ToTensord(keys=["image", "label"]),
        ]
    )
    val_transform = Compose(
        [
            LoadImaged(keys=["image", "label"]),
            AddChanneld(keys=["image"]),
            ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
            Spacingd(
                keys=["image", "label"],
                pixdim=(1,1,1),
                mode=("bilinear", "nearest"),
            ),
            Orientationd(
                keys=["image", "label"],
                axcodes="LPS"
            ),
            ResizeWithPadOrCropd(keys=["image", "label"],
                spatial_size=(256,256,256),
                mode="constant"),
            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
            ToTensord(keys=["image", "label"]),
        ]
    )

Establish training datasets

    train_ds = CacheDataset(
        data=train_files,
        transform=train_transform,
        cache_rate=0.5,
        num_workers=4)
    train_loader = DataLoader(
        train_ds,
        batch_size=1,
        shuffle=True,
        num_workers=4) 
    val_ds = CacheDataset(
        data=val_files,
        transform=val_transform,
        cache_rate=0.5,
        num_workers=4)
    val_loader = DataLoader(
        val_ds,
        batch_size=1,
        shuffle=False,
        num_workers=4) 
    
    print(f"train image shape: {train_ds[12]['image'].shape}")
    print(f"train label shape: {train_ds[12]['label'].shape}")
    
    print(f"val image shape: {val_ds[4]['image'].shape}")
    print(f"val label shape: {val_ds[4]['label'].shape}")

Visualize input image:

    print(f"image shape: {train_ds[15]['image'].shape}")
    plt.figure("Train image", (14, 6))
    for i in range(1):
        plt.subplot(1, 3, i + 1)
        plt.title(f"image channel {i}")
        plt.imshow(train_ds[15]["image"][i, :, :, 120].detach().cpu(), cmap="gray")
    plt.savefig('/users/jmenvill/ICH/Models/Visualize_Training_Image_15.png', bbox_inches='tight')

Visualize corresponding label channels:

    print(f"label shape: {train_ds[15]['label'].shape}")
    plt.figure("Train label", (14, 6))
    for i in range(2):
        plt.subplot(1, 2, i + 1)
        plt.title(f"label channel {i}")
        plt.imshow(train_ds[15]["label"][i, :, :, 120].detach().cpu())
    plt.savefig('/users/jmenvill/ICH/Models/Visualize_Training_Labels_15.png', bbox_inches='tight')

standard PyTorch program style: create UNet, DiceLoss and Adam optimizer

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = UNet(
        dimensions=3,
        in_channels=1,
        out_channels=2,
        channels=(16, 32, 64, 128, 256),
        strides=(2, 2, 2, 2),
        num_res_units=2,
    ).to(device)
    loss_function = DiceLoss(to_onehot_y=False, sigmoid=True, squared_pred=True)
    optimizer = torch.optim.Adam(
        model.parameters(), 1e-4, weight_decay=1e-5, amsgrad=True
    )
    dice_metric = DiceMetric(include_background=True, reduction="mean")
    dice_metric_batch = DiceMetric(include_background=True, reduction="mean_batch")
    post_trans = Compose(
        [ToTensor(), Activations(sigmoid=True), AsDiscrete(threshold_values=True)]
    )

Execute training

    max_epochs = 250
    val_interval = 2
    best_metric = -1
    best_metric_epoch = -1
    epoch_loss_values = []
    metric_values = []
    metric_values_he = []
    metric_values_ed = []
    for epoch in range(max_epochs):
        print("-" * 10)
        print(f"epoch {epoch + 1}/{max_epochs}")
        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()
            print(
                f"{step}/{len(train_ds) // train_loader.batch_size}"
                f", 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():
    
                for val_data in val_loader:
                    val_inputs, val_labels = (
                        val_data["image"].to(device),
                        val_data["label"].to(device),
                    )
                    val_outputs = model(val_inputs)
                    val_outputs = [post_trans(i) for i in decollate_batch(val_outputs)]
                    dice_metric(y_pred=val_outputs, y=val_labels)
                    dice_metric_batch(y_pred=val_outputs, y=val_labels)
    
                metric = dice_metric.aggregate().item()
                metric_values.append(metric)
                metric_batch = dice_metric_batch.aggregate()
                
                metric_he = metric_batch[0].item()
                metric_values_he.append(metric_he)
                metric_ed = metric_batch[1].item()
                metric_values_ed.append(metric_ed)
                
                dice_metric.reset()
                dice_metric_batch.reset()
    
                if metric > best_metric:
                    best_metric = metric
                    best_metric_epoch = epoch + 1
                    torch.save(
                        model.state_dict(),
                        os.path.join(model_dir, "best_model_uniform_sizing.pth"),
                    )
                    print()
                    print("saved new best metric model")
                
                print(
                    f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                    f" HEMATOMA: {metric_he:.4f} EDEMA: {metric_ed:.4f}"
                    f"\nbest mean dice: {best_metric:.4f}"
                    f" at epoch: {best_metric_epoch}"
                )

Print best metric

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

Plot the loss and metric

    plt.figure("train", (12, 6))
    plt.subplot(1, 2, 1)
    plt.title("Epoch Average Loss")
    x = [i + 1 for i in range(len(epoch_loss_values))]
    y = epoch_loss_values
    plt.xlabel("epoch")
    plt.plot(x, y, color="red")
    plt.savefig('/users/jmenvill/ICH/Models/Epoch_Average_Loss.png', bbox_inches='tight')
    
    plt.subplot(1, 2, 2)
    plt.title("Val Mean Dice")
    x = [val_interval * (i + 1) for i in range(len(metric_values))]
    y = metric_values
    plt.xlabel("epoch")
    plt.plot(x, y, color="orange")
    plt.show()
    plt.savefig('/users/jmenvill/ICH/Models/Overall_Val_Mean_Dice.png', bbox_inches='tight')
    
    plt.figure("train", (18, 6))
    plt.subplot(1, 3, 1)
    plt.title("Val Mean Dice HEMATOMA")
    x = [val_interval * (i + 1) for i in range(len(metric_values_he))]
    y = metric_values_he
    plt.xlabel("epoch")
    plt.plot(x, y, color="blue")
    plt.savefig('/users/jmenvill/ICH/Models/Hematoma_Val_Mean_Dice.png', bbox_inches='tight')
    
    plt.subplot(1, 3, 2)
    plt.title("Val Mean Dice EDEMA")
    x = [val_interval * (i + 1) for i in range(len(metric_values_ed))]
    y = metric_values_ed
    plt.xlabel("epoch")
    plt.plot(x, y, color="purple")
    plt.show()
    plt.savefig('/users/jmenvill/ICH/Models/Edema_Val_Mean_Dice.png', bbox_inches='tight')

Check best model output with the input image and label

    model.load_state_dict(
        torch.load(os.path.join(model_dir, "best_model_uniform_sizing.pth"))
    )
    model.eval()
    with torch.no_grad():
        
        # select one image to evaluate and visualize the model output
        val_input = val_ds[0]["image"].unsqueeze(0).to(device)
        val_output = model(val_input)
        plt.figure("image", (24, 6))
        for i in range(1):
            plt.subplot(1, 1, i+1)
            plt.title(f"image channel = FLAIR")
            plt.imshow(val_ds[0]["image"][i, :, :, 120].detach().cpu(), cmap="gray")
        plt.savefig('/users/jmenvill/ICH/Models/Visualize_Model_Image_Output_0.png', bbox_inches='tight')

visualize the 2 channels label corresponding to this image

        plt.figure("label", (24, 6))
        for i in range(2):
            plt.subplot(1, 2, i + 1)
            plt.title(f"label channel {i}")
            plt.imshow(val_ds[0]["label"][i, :, :, 120].detach().cpu())
        plt.savefig('/users/jmenvill/ICH/Models/Visualize_Model_Labels_Output_0.png', bbox_inches='tight')

visualize the 2 channels model output corresponding to this image

        plt.figure("output", (24, 6))
        for i in range(2):
            plt.subplot(1, 2, i + 1)
            plt.title(f"output channel {i}")
            plt.imshow(val_output[0, i, :, :, 120].detach().cpu())
        plt.savefig('/users/jmenvill/ICH/Models/Visualize_Model_Channels_Output_0.png', bbox_inches='tight')

[Nic-Ma]

Hi @jmenville11 ,

Thanks for your interest and experiments with MONAI.
Could you please help share the full error log for further analysis?
@IsaacYangSLA I remember you encountered some similar issue before when you were training brats, could you please help share some comments or advice here?

Thanks in advance.

[Nic-Ma]

Seems like some PyTorch IO reading issue with your storage.
Could you please help try again with num_workers=0 in your DataLoader?

Thanks.

[jmenville11]

Thanks for the suggestion, the model ran with a Cache Rate of 0.5 when num_workers = 0. I still don't fully understand how the parameters are interacting, though. Is it best to always keep num_workers = 0?

[Nic-Ma]

Hi @jmenville11 ,

I think it's mainly related to the PyTorch DataLoader, especially the multi-processing IO, etc.
https://github.com/pytorch/pytorch/blob/master/torch/utils/data/dataloader.py#L693

Thanks.

[jmenville11]

I see. I’ve read through the linked doc a few times now but don’t see a strong explanation that explains the interaction. The doc doesn’t reference anything about cache rate.

[Nic-Ma]

Yes, if you cached many things in the memory, it may easier to cause some multi-processing communication issue I think.

Thanks.