Transformer UNET will not train even with diverse data and augmentations

theskywalker: <p>Hi everyone. I am building a UNETR, which is a Transformer based UNET for medical image vertebrae instance segmentation with MONAI. I am using the VerSe dataset, which is high quality and diverse CT images with labels of the human spine.</p>
<p>Here is my code where I perform data augmentation on the images. I use scaleintensityranged() to threshold only the brightest bone structures (like the spine)</p>
<blockquote>
<p>train_transforms = Compose(<br>
[<br>
LoadImaged(keys=[“image”, “label”]),<br>
EnsureChannelFirstd(keys=[“image”, “label”]),<br>
Orientationd(keys=[“image”, “label”], axcodes=“RAS”),<br>
Spacingd(<br>
keys=[“image”, “label”],<br>
pixdim=(1.17, 1.17, 5.0),<br>
mode=(“bilinear”, “nearest”),<br>
),<br>
ScaleIntensityRanged(<br>
keys=[“image”],<br>
a_min=200,<br>
a_max=400,<br>
b_min=0.0,<br>
b_max=1.0,<br>
clip=True,<br>
),<br>
CropForegroundd(keys=[“image”, “label”], source_key=“image”,allow_smaller=True),<br>
RandCropByPosNegLabeld(<br>
keys=[“image”, “label”],<br>
label_key=“label”,<br>
spatial_size=(96, 96, 96),<br>
pos=1,<br>
neg=1,<br>
num_samples=4,<br>
image_key=“image”,<br>
image_threshold=0,<br>
),<br>
RandFlipd(<br>
keys=[“image”, “label”],<br>
spatial_axis=[0],<br>
prob=0.10,<br>
),<br>
RandFlipd(<br>
keys=[“image”, “label”],<br>
spatial_axis=[1],<br>
prob=0.10,<br>
),<br>
RandFlipd(<br>
keys=[“image”, “label”],<br>
spatial_axis=[2],<br>
prob=0.10,<br>
),<br>
RandRotate90d(<br>
keys=[“image”, “label”],<br>
prob=0.10,<br>
max_k=3,<br>
),<br>
RandShiftIntensityd(<br>
keys=[“image”],<br>
offsets=0.10,<br>
prob=0.50,<br>
),<br>
]<br>
)<br>
)<br>
val_transforms = Compose(<br>
[<br>
LoadImaged(keys=[“image”, “label”]),<br>
EnsureChannelFirstd(keys=[“image”, “label”]),<br>
Orientationd(keys=[“image”, “label”], axcodes=“RAS”),<br>
Spacingd(<br>
keys=[“image”, “label”],<br>
pixdim=(1.17, 1.17, 5.0),<br>
mode=(“bilinear”, “nearest”),<br>
),<br>
ScaleIntensityRanged(keys=[“image”], a_min=200, a_max=400, b_min=0.0, b_max=1.0, clip=True),<br>
CropForegroundd(keys=[“image”, “label”], source_key=“image”,allow_smaller=True),<br>
]<br>
)</p>
</blockquote>
<p>This is how I prepared the data. As you can see, the batch size is 2.</p>
<blockquote>
<p>train_ds = CacheDataset(<br>
data=train_files,<br>
transform=train_transforms,<br>
cache_num=24,<br>
cache_rate=1.0,<br>
num_workers=8,<br>
)<br>
batch_size = 2<br>
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True)<br>
val_ds = CacheDataset(data=val_files, transform=val_transforms, cache_num=6, cache_rate=1.0, num_workers=4)<br>
val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)</p>
</blockquote>
<p>Below is my model, loss function, optimizer, and training code:</p>
<blockquote>
<p>os.environ[“CUDA_DEVICE_ORDER”] = “PCI_BUS_ID”<br>
CUDA_LAUNCH_BLOCKING=1<br>
torch.backends.cudnn.benchmark = True<br>
torch.set_float32_matmul_precision(‘medium’)<br>
no_classes = 29<br>
model = UNETR(<br>
in_channels=1,<br>
out_channels=no_classes,<br>
img_size=(96, 96, 96),<br>
feature_size=16,<br>
hidden_size=768,<br>
mlp_dim=3072,<br>
num_heads=12,<br>
proj_type=“perceptron”,<br>
norm_name=“batch”,<br>
res_block=True,<br>
conv_block=True,<br>
dropout_rate=0,<br>
).to(device)</p>
</blockquote>
<blockquote>
<p>lr = 1e-3<br>
max_iterations = 6000<br>
eval_num = 100<br>
post_label = AsDiscrete(to_onehot=no_classes)<br>
post_pred = AsDiscrete(argmax=True, to_onehot=no_classes)<br>
dice_metric = DiceMetric(include_background=False, reduction=“mean”, get_not_nans=False)<br>
loss_function = DiceLoss(to_onehot_y=True, softmax=True,include_background=False)<br>
optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)<br>
scaler = torch.cuda.amp.GradScaler()<br>
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100,1000], gamma=0.1)</p>
</blockquote>
<blockquote>
<p>def validation(epoch_iterator_val):<br>
model.eval()<br>
with torch.no_grad():<br>
for batch in epoch_iterator_val:<br>
val_inputs, val_labels = (batch[“image”].to(device), batch[“label”].to(device))<br>
with torch.cuda.amp.autocast():<br>
val_outputs = sliding_window_inference(val_inputs, (96, 96, 96), 4, model)<br>
val_labels_list = decollate_batch(val_labels)<br>
val_labels_convert = [post_label(val_label_tensor) for val_label_tensor in val_labels_list]<br>
val_outputs_list = decollate_batch(val_outputs)<br>
val_output_convert = [post_pred(val_pred_tensor) for val_pred_tensor in val_outputs_list]<br>
dice_metric(y_pred=val_output_convert, y=val_labels_convert)<br>
epoch_iterator_val.set_description(“Validate (%d / %d Steps)” % (global_step, 10.0))  # noqa: B038<br>
mean_dice_val = dice_metric.aggregate().item()<br>
dice_metric.reset()<br>
return mean_dice_val</p>
</blockquote>
<blockquote>
<p>def train(global_step, train_loader, dice_val_best, global_step_best):<br>
model.train()<br>
epoch_loss = 0<br>
step = 0<br>
epoch_iterator = tqdm(train_loader, desc=“Training (X / X Steps) (loss=X.X)”, dynamic_ncols=True)<br>
for step, batch in enumerate(epoch_iterator):<br>
step += 1<br>
x, y = (batch[“image”].to(device), batch[“label”].to(device))<br>
with torch.cuda.amp.autocast():<br>
logit_map = model(x)<br>
loss = loss_function(logit_map, y)<br>
scaler.scale(loss).backward()<br>
epoch_loss += loss.item()<br>
scaler.unscale_(optimizer)<br>
scaler.step(optimizer)<br>
scale = scaler.get_scale()<br>
scaler.update()<br>
# if not (scale &gt; scaler.get_scale()):<br>
scheduler.step()<br>
optimizer.zero_grad()<br>
epoch_iterator.set_description(  # noqa: B038<br>
“Training (%d / %d Steps) (loss=%2.5f)” % (global_step, max_iterations, loss)<br>
)<br>
if (global_step % eval_num == 0 and global_step != 0) or global_step == max_iterations:<br>
epoch_iterator_val = tqdm(val_loader, desc=“Validate (X / X Steps) (dice=X.X)”, dynamic_ncols=True)<br>
dice_val = validation(epoch_iterator_val)<br>
epoch_loss /= step<br>
epoch_loss_values.append(epoch_loss)<br>
metric_values.append(dice_val)<br>
if dice_val &gt; dice_val_best:<br>
dice_val_best = dice_val<br>
global_step_best = global_step<br>
torch.save({<br>
‘model_state_dict’: model.state_dict(),<br>
‘optim_state_dict’: optimizer.state_dict(),<br>
‘loss’: loss,<br>
},os.path.join(root_dir, “best_metric_model.ckpt”))<br>
print(<br>
“Model Was Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}”.format(dice_val_best, dice_val)<br>
)<br>
else:<br>
print(<br>
“Model Was Not Saved ! Current Best Avg. Dice: {} Current Avg. Dice: {}”.format(<br>
dice_val_best, dice_val<br>
)<br>
)<br>
global_step += 1<br>
print("Current Learning Rate: ", scheduler.get_last_lr()[0])<br>
print("Epoch Loss: ", epoch_loss)<br>
return global_step, dice_val_best, global_step_best</p>
</blockquote>
<p>Finally, here is an example of the output of training. Even after 8+ hours, the model stays exactly the same:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/8/884a9b071e4913aa64fdd293d390949b9b6ee2a8.png" data-download-href="https://discuss.pytorch.org/uploads/default/884a9b071e4913aa64fdd293d390949b9b6ee2a8" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/8/884a9b071e4913aa64fdd293d390949b9b6ee2a8.png" alt="image" data-base62-sha1="jrGMds6WqJzNwg9PIMuP4CtipKw" width="652" height="500" data-dominant-color="303030"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">824×631 28.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
Any help would be so so appreciated. I wanted to post more images of the data and augmentation for extra clarity but it wouldn’t let me!</p>

Aniruth_Sundararaja1: <p>What is your batch size , keeping a very low batch size is bad.<br>
It will not let your model learn effectively.<br>
Try to keep a minimum batch size of 8.</p>
<p>If you could share your model architecture , it would be easier to find the issue.</p>

theskywalker: <p>Hi Aniruth,<br>
I edited my post to add how I prepared my date using the Cache Dataloader.<br>
My training batch size is 2. I experimented quite a bit with larger batch sizes but 2 was the largest I could get without my GPU giving me memory errors. I have an RTX 4070 with 12 GB of VRAM, but since the UNETR is a 3D model, it uses a ton of memory.</p>

Aniruth_Sundararaja1: <p>Did you try converting your image into patches and train your model with the patches instead of a single high resolution image. That helps for high resolution images.</p>
<p>In that case you can increase your batch size too</p>