Stack expects each tensor to be equal size, but got [163, 256, 256] at entry 0 and [160, 256, 256] at entry 1

vision
1

I am working with the OAI MRI dataset for knee osteoarthritis classification. Each one of 435 MRIs I got has to be classified to a grade. For each MRI in a folder, there are 160 2D images. I created this function to read the dataset:

def dicom2array(path):
  dicom = pydicom.read_file(path)
  data = dicom.pixel_array
  data = (data - np.min(data)) / (np.max(data) - np.min(data))
  data = cv2.resize(data, (256,256))
  return data

def load_3d_dicom_images(scan_id): #returns an object with shape (160,256,256)
  files = sorted(glob.glob(f"/content/drive/MyDrive/subjects/{scan_id}/*/*"))
  img = np.array([dicom2array(a) for a in files])
  return img

class MyCustomDataset(Dataset):
  def __init__(self, path = "/content/drive/MyDrive/subjects"):
    dataframe = pd.read_csv("/content/drive/MyDrive/KL_grade.csv", sep = ';')
    self.labels = {}
    id = list(dataframe["id"])
    grades = list(dataframe["grade"])
    for i,g in zip(id, grades):
      self.labels[str(i).zfill(5)] = g
    
    self.ids = [a.split("/")[-1] for a in sorted(glob.glob(f"/content/drive/MyDrive/subjects/" + "/*"))]

  def __len__(self):
    return len(self.ids)

  def __getitem__(self, idx):
    imgs = load_3d_dicom_images(self.ids[idx])  
    label = self.labels[self.ids[idx]]

    return torch.tensor(imgs, dtype = torch.float32), torch.tensor(label, dtype = torch.long)

I tried to train the model with a resnet18 from the monai library. When I set the batch_size=10 or more cuda runs out of memory (I run it on google colab). It works with batch_size=1 but when I set it equal to 2 I get this error:
stack expects each tensor to be equal size, but got [163, 256, 256] at entry 0 and [160, 256, 256] at entry 1
Can anybody help me ?

2

Based on the error message some samples seem to have 160 channels while others have 163.
Due to this, the collate_fn cannot stack the samples to a single batch and raises the error.
Check the shape of each sample and make sure the channel dimension as well as the spatial size are the same.

3

Thanks for your response!
If that was the reason the error occured, why the model ran with batch size = 1 ?

4

If you are using a batch size of 1, the collate_fn will “stack” the single sample to a batch, i.e. it will just add the batch dimension to the sample and return it without trying to stack several samples to a single batch.

5

Ok I think I got it, thank you!
I have another more general question if you have time to answer.
The class I wrote to get the dataset do you think it’s ok ?

6

Your custom Dataset looks good as it lazily loads and processes the images.

7

Is there any way that I can make it not beeing lazy ?

8

Yes, you can preload the data in the __init__ method, but why would you want to do it?
Often you don’t want to preload the data due to the potentially large memory requirement and thus slower initialization. Also, you would use multiprocessing to load and process each batch (lazily) in the background while the model trains, so lazy loading is usually the preferred option.

9

Is there any problem for the general perfomance of my model because of the low batch size ?

10

Increasing the batch size would also increase the overall device utilization assuming you are not creating a data loading bottleneck, but since I’m not familiar with your model I don’t know how much improvement an increase in the batch size would give. You could profile your code and check the improvement via the native PyTorch profiler or e.g. Nsight Systems.

11

I am asking because I am limited from google colab to set batch size equal to 2. If I set it any larger than 2 I get error that CUDA runs out of memory