way to improve the output

[pure-rgb]

I'm using Covid-19 CT segmentaiton data set,

and I am following this set up https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/spleen_segmentation_3d.ipynb with some necessary changes.

The data contains 4 labels, ie. background, left lungs, right lungs and covid symtoms - data set links. Now, even I trained the model with 500 / 600 epochs, I couldn't get better results, dice mean score only around 0.60. Below is the code. Could you please suggest what things should I consider? Do I need to train much longer with more epochs (1000 / 2000). Or, is there (below code) config I set wrongly. Should I consider it multi-class or multi-label? What changes in code I need to get good score? Thank you.

data_dir = './covid'
train_images = sorted(
    glob.glob(os.path.join(data_dir, "COVID-19-CT-Seg_20cases", "*.nii")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "Lung_and_Infection_Mask", "*.nii")))

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

num_samples = 4

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"], 
            a_min=-57, 
            a_max=164,
            b_min=0.0, 
            b_max=1.0, 
            clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"], 
            pixdim=(1.5, 1.5, 2.0), 
            mode=("bilinear", "nearest")
        ),
        RandCropByPosNegLabeld(
            keys=["image", "label"],
            label_key="label",
            spatial_size=(96, 96, 96),
            pos=1,
            neg=1,
            num_samples=4,
            image_key="image",
            image_threshold=0,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[0],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[1],
            prob=0.10,
        ),
        RandFlipd(
            keys=["image", "label"],
            spatial_axis=[2],
            prob=0.10,
        ),
        RandRotate90d(
            keys=["image", "label"],
            prob=0.10,
            max_k=3,
        ),
        RandShiftIntensityd(
            keys=["image"],
            offsets=0.10,
            prob=0.50,
        ),
    ]
)
val_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        ScaleIntensityRanged(
            keys=["image"], a_min=-57, a_max=164,
            b_min=0.0, b_max=1.0, clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image"),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"], 
            pixdim=(1.5, 1.5, 2.0), 
            mode=("bilinear", "nearest")
        ),
    ]
)


train_ds = CacheDataset(
    data=train_files, 
    transform=train_transforms,
    cache_rate=1.0, 
    num_workers=4
)
train_loader = DataLoader(
    train_ds, 
    batch_size=1, 
    shuffle=True, 
    num_workers=4
)

val_ds = CacheDataset(
    data=val_files, 
    transform=val_transforms, 
    cache_rate=1.0, 
    num_workers=4
)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
device = torch.device("cuda:0")


model = SwinUNETR(
    img_size=(96, 96, 96),
    in_channels=1,
    out_channels=4,
    feature_size=24,
    drop_rate=0.0,
    attn_drop_rate=0.0,
    dropout_path_rate=0.0,
    use_checkpoint=True,
).to(device)

# model = UNet(
#     spatial_dims=3,
#     in_channels=1,
#     out_channels=4,
#     channels=(16, 32, 64, 128, 256),
#     strides=(2, 2, 2, 2),
#     num_res_units=2,
#     norm=Norm.BATCH,
# ).to(device)

# torch.backends.cudnn.benchmark = True
loss_function = DiceLoss(to_onehot_y=True, softmax=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=True, reduction="mean")

post_pred = Compose([AsDiscrete(argmax=True, to_onehot=4)])
post_label = Compose([AsDiscrete(to_onehot=4)])

max_epochs = 2000
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []

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),
                )
                
                roi_size = (96, 96, 96)
                # roi_size = (160, 160, 160)
                
                sw_batch_size = 4
                val_outputs = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                # 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(), 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} "
                f"at epoch: {best_metric_epoch}"
            )

[ebrahimebrahim]

Do I need to train much longer with more epochs (1000 / 2000). Or, is there (below code) config I set wrongly. Should I consider it multi-class or multi-label? What changes in code I need to get good score? Thank you.

How many epochs you need depends very much on the situation and the learning rate. Inspect your loss curves to see whether more epochs would do more good. If the training loss looks like it could still decrease with more epochs, then more training could help. But also keep watching the validation loss curve; if validation loss increases then you are beginning to overfit and will not gain generalizable performance with more training (and might even lose perforamance).

The data contains 4 labels, ie. background, left lungs, right lungs and covid symtoms - data set links. Now, even I trained the model with 500 / 600 epochs, I couldn't get better results, dice mean score only around 0.60. Below is the code. Could you please suggest what things should I consider?

Maybe it's not helpful to treat "covid symptoms" as a separate segmentation class. A voxel in the left lung that shows covid symptoms is still a "left lung" voxel, and it's weird to penalize the model for the high probability it assigns to "left lung" for that voxel.

One idea: You could split things up so that the lung segmentation is separate from finding covid symptoms. For example the output channels could be treated like this: One channel is the class probability (or log-probability) for being lung rather than background, and another channel is class probability for being part of a covid symptom region rather than not being part of a covid symptom region. (Notice these two channels defined this way would not be expected to sum to 1.) Then you could have two separate dice losses, one for each of the two problems: lung segmentation and covid symptom localization. This assumes your data is such that "covid symptom" always occurs within the lungs, so that taking the union of the covid symptom and the lung labels provides you with a complete lung label.

[pure-rgb]

Thanks for the suggestion. I'm new to medical imaging. I train the model much longer, ~2000 epochs. The mean dice score achive ~0.73. I think the above coid-19 dataset is quite similar with this example, I tried to follow this one too but very disapointing results. https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/swin_unetr_btcv_segmentation_3d.ipynb

[Nic-Ma]

Hi @dongyang0122 ,

Could you please help share some best practices for the training?

Thanks in advance.

[dongyang0122]

hi @Suzan009, could you please share your expectation of the average (validation) Dice score? As @ebrahimebrahim suggested, it would be helpful if you can share the training curves (e.g., training loss, validation Dice score). And I noticed two places in your scripts that might affect the performance.

1. The window [-57, 164] in ScaleIntensityRanged is too small for lung regions. Some appearance cues of COVID-19 lesions in CT images might be reduced after intensity normalization. You can use the recommendation window in this link or this link.
2. The re-sampling spacing (1.5, 1.5, 2.0) might limit the performance improvement because some details of labels would be smoothed after re-sampling. You can increase the spacing to ((1.5, 1.5, 1.5) or (1.0, 1.0, 1.0) for further improvement.