UNet trains, UNETR does not on same data and augmentations

[theskywalker1]

Hello Monai community.
I am attempting to train a UNETR model using this tutorial for vertebral segmentation.
Using this tutorial, I was able to sucessfully train a regular UNet on the VerSe dataset. I chose 25 high quality, full body volumes with a 70/30 split for training and validation.
Here is an example sagittal slice of this dataset:

Here is my config for this project:

In the code snippets I provide below, you will see model1 as a UNETR model and model2 as the regular UNet.
Below this, I provide a graph of the regular UNet's performance and a screenshot of the UNETR's lack of change.
Could anyone explain why a more complex and supposedly better performing model is behaving much worse on the same dataset?
Thanks for anything.

train_transforms = Compose(
    [
        LoadImaged(keys=["image", "label"]),
        EnsureChannelFirstd(keys=["image", "label"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.17, 1.17, 5.0),
            mode=("bilinear", "nearest"),
        ),
        ScaleIntensityRanged(
            keys=["image"],
            a_min=-200,
            a_max=164,
            b_min=0.0,
            b_max=1.0,
            clip=True,
        ),
        CropForegroundd(keys=["image", "label"], source_key="image",allow_smaller=True),
        # EnsureTyped(keys=["image", "label"],device=device,track_meta=False),
        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"]),
        Orientationd(keys=["image", "label"], axcodes="RAS"),
        Spacingd(
            keys=["image", "label"],
            pixdim=(1.17, 1.17, 5.0),
            mode=("bilinear", "nearest"),
        ),
        ScaleIntensityRanged(keys=["image"], a_min=-200, a_max=164, b_min=0.0, b_max=1.0, clip=True),
        CropForegroundd(keys=["image", "label"], source_key="image",allow_smaller=True),
        # EnsureTyped(keys=["image", "label"], device=device, track_meta=True),
    ]
)

train_ds = CacheDataset(
    data=train_files,
    transform=train_transforms,
    cache_num=24,
    cache_rate=1.0,
    num_workers=8,
)
batch_size = 2
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
CUDA_LAUNCH_BLOCKING=1
torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('medium')
no_classes = 29
ReTrain = False
model1 = UNETR(
    in_channels=1,
    out_channels=no_classes,
    img_size=(96, 96, 96),
    feature_size=16,
    hidden_size=768,
    mlp_dim=3072,
    num_heads=12,
    proj_type="perceptron",
    norm_name="batch",
    res_block=True,
    conv_block=True,
    dropout_rate=0,
).to(device)
model2 = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=no_classes,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)

lr = 1e-4
max_iterations = 6000 
eval_num = 100

post_label = AsDiscrete(to_onehot=no_classes)
post_pred = AsDiscrete(argmax=True, to_onehot=no_classes)
dice_metric = DiceMetric(include_background=False, reduction="mean", get_not_nans=False)

loss_function = DiceLoss(to_onehot_y=True, softmax=True,include_background=False)
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
scaler = torch.cuda.amp.GradScaler()
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', patience=eval_num, factor=0.5)
if ReTrain == True:
    checkpoint = torch.load(os.path.join(root_dir, "best_metric_model.ckpt"))
    model.load_state_dict(checkpoint["model_state_dict"])
    optimizer.load_state_dict(checkpoint["optim_state_dict"])
    loss = checkpoint["loss"]

def validation(epoch_iterator_val):
    model.eval()
    with torch.no_grad():
        for batch in epoch_iterator_val:
            val_inputs, val_labels = (batch["image"].to(device), batch["label"].to(device))
            with torch.cuda.amp.autocast():
                val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model)
            val_labels_list = decollate_batch(val_labels)
            val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]
            val_outputs_list = decollate_batch(val_outputs)
            val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]
            dice_metric(y_pred=val_output_convert, y=val_labels_convert)
            epoch_iterator_val.set_description("Validate (%d / %d Steps)" % (global_step, 10.0))  # noqa: B038
        mean_dice_val = dice_metric.aggregate().item()
        dice_metric.reset()
    return mean_dice_val

def train(global_step, train_loader, dice_val_best, global_step_best):
    model.train()
    epoch_loss = 0
    step = 0
    epoch_iterator = tqdm(train_loader, desc="Training (X / X Steps) (loss=X.X)", dynamic_ncols=True)
    for step, batch in enumerate(epoch_iterator):
        step += 1
        x, y = (batch["image"].to(device), batch["label"].to(device))
        with torch.cuda.amp.autocast():
            logit_map = model(x)
            loss = loss_function(logit_map, y)
        scaler.scale(loss).backward()
        epoch_loss += loss.item()
        scaler.unscale_(optimizer)
        scaler.step(optimizer)
        scale = scaler.get_scale()
        scaler.update()
        # if not (scale > scaler.get_scale()):
        scheduler.step(loss)
        optimizer.zero_grad()
        epoch_iterator.set_description(  # noqa: B038
            "Training (%d / %d Steps) (loss=%2.5f)" % (global_step, max_iterations, loss)
        )
        if (global_step % eval_num == 0 and global_step != 0) or global_step == max_iterations:
            epoch_iterator_val = tqdm(val_loader, desc="Validate (X / X Steps) (dice=X.X)", dynamic_ncols=True)
            dice_val = validation(epoch_iterator_val)
            epoch_loss /= step
            epoch_loss_values.append(epoch_loss)
            metric_values.append(dice_val)
            if dice_val > dice_val_best:
                dice_val_best = dice_val
                global_step_best = global_step
                torch.save({
                    'model_state_dict': model.state_dict(),
                    'optim_state_dict': optimizer.state_dict(),
                    'loss': loss,
                    },os.path.join(root_dir, "best_metric_model1.ckpt"))
                print(
                    "Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(dice_val_best, dice_val)
                )
            else:
                print(
                    "Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}".format(
                        dice_val_best, dice_val
                    )
                )
        global_step += 1
    print("Current learning rate: ", optimizer.param_groups[0]["lr"])
    print("Epoch Loss: ", epoch_loss)
    return global_step, dice_val_best, global_step_best

global_step = 0
dice_val_best = 0.0
global_step_best = 0
epoch_loss_values = []
metric_values = []
start = datetime.now()
while global_step < max_iterations:
    global_step, dice_val_best, global_step_best = train(global_step, train_loader, dice_val_best, global_step_best)
end=datetime.now()-start
print('Time to train:',end)

Here is the performance over 6k iterations on the regular UNet:

Unfortunately, I do not have a graph of the UNETR's training. But I do have this screenshot of the console showing no change in the loss or mean dice:

[KumoLiu]

Hi @theskywalker1, this is not to say that a more complex network will always yield better results; in your example, unet can only achieve an accuracy of roughly 0.4. When switching to a more complex UNETR, you should adjust the parameters appropriately. Even so, it frequently happens that we overfit a basic network before attempting to intentionally alter it.

Thanks.