UNet 3D and LearningRateFinder

[pauljhin]

Hello,

I'm trying to use LearningRateFinder like monai tutorial [1] but with spleen dataset and UNet (3D dimension) and this error is occurring:

Computing optimal learning rate:   0%|                                                                                                                                                      | 0/20 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "teste_find_lr.py", line 62, in <module>
    lr_finder.range_test(train_loader, val_loader, end_lr=upper_lr, num_iter=20)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/optimizers/lr_finder.py", line 343, in range_test
    loss = self._train_batch(train_iter, accumulation_steps, non_blocking_transfer=non_blocking_transfer)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/optimizers/lr_finder.py", line 400, in _train_batch
    outputs = self.model(inputs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/networks/nets/unet.py", line 311, in forward
    x = self.model(x)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/networks/layers/simplelayers.py", line 120, in forward
    y = self.submodule(x)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/networks/layers/simplelayers.py", line 120, in forward
    y = self.submodule(x)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/container.py", line 139, in forward
    input = module(input)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/networks/layers/simplelayers.py", line 123, in forward
    return torch.cat([x, y], dim=self.dim)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/monai/data/meta_tensor.py", line 249, in __torch_function__
    ret = super().__torch_function__(func, types, args, kwargs)
  File "/home/usr/.virtualenvs/ai/lib/python3.8/site-packages/torch/_tensor.py", line 1121, in __torch_function__
    ret = func(*args, **kwargs)
  RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 3 but got size 4 for tensor number 1 in the list.

My code:

import glob 
import os
import numpy as np 
import torch 
from monai.data import *
from monai.transforms import *
import torch.nn as nn
from monai.optimizers import LearningRateFinder
from monai.metrics import DiceMetric
from monai.losses import DiceLoss
from monai.networks.nets import UNet
from monai.networks.layers import Norm

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

transforms = Compose(
    [
        LoadImaged(keys=["image","label"]),
        EnsureChannelFirstd(keys=["image","label"]),
        ScaleIntensityd(keys=["image","label"]),
        CenterSpatialCropd(keys=["image","label"], roi_size=(20, 20,20)),
        EnsureTyped(keys=["image","label"]),
    ]
)

data_dir="./mydata/"

train_images = sorted(
    glob.glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
train_labels = sorted(
    glob.glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
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:]

train_ds = CacheDataset(data=train_files, transform=transforms,cache_rate=1, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=1, num_workers=4)

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

device="cuda"
model = model.to(device)
loss_function = DiceLoss(to_onehot_y=True, softmax=True)

lower_lr, upper_lr = 1e-3, 1e-0
optimizer = torch.optim.Adam(model.parameters(), lower_lr)
lr_finder = LearningRateFinder(model, optimizer, loss_function, device=device)
lr_finder.range_test(train_loader, val_loader, end_lr=upper_lr, num_iter=20)
steepest_lr, _ = lr_finder.get_steepest_gradient()
print(steepest_lr)

[1] - https://github.com/Project-MONAI/tutorials/blob/main/modules/learning_rate.ipynb

Thanks in advance!

[ericspod]

The spatial dimensions for the input to a UNet with this setup have to be multiples of 2**4 since you're downsampling by half in spatial dimensions four times as defined in the strides parameter. Do you know what the dimensions are the input files? You can use one of the crop or pad transforms to make the sizes fit, eg. DivisiblePadd.

[pauljhin]

@ericspod the dimension is 512 x 512 x N. It worked with your help. Thank's a lot!