Encountered Dimension Mismatch Error While Training DRENet

[wenjieo]

Hello, I encountered an issue while training the DRENet model with the SSDD dataset. Due to the inconsistent image dimensions in SSDD, the torch.stack operation in the collate_fn function fails, preventing normal training. Could you please provide any suggestions on how to best handle the image resizing or any other solutions that might work well with your model setup? The error message is as follows:
Image sizes 512 train, 512 test
Using 8 dataloader workers
Logging results to run\train\exp16
Starting training for 300 epochs...

 Epoch   gpu_mem       box       obj       cls       dgi     total   targets  img_size
 0/299     4.43G    0.1236   0.02271         0   0.02122     2.856        26       512: 100%|██████████| 93/93 [00:15<00:00,  5.88it/s]
           Class      Images     Targets           P           R      mAP@.5  mAP@.5:.95:   0%|          | 0/24 [00:00<?, ?it/s]

Traceback (most recent call last):
File "E:\ship detection\DRENet-main\train.py", line 521, in
train(hyp, opt, device, tb_writer, wandb)
File "E:\ship detection\DRENet-main\train.py", line 352, in train
results, maps, times = test.test(opt.data,
File "E:\ship detection\DRENet-main\test.py", line 103, in test
for batch_i, (img, dgimgs, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\tqdm\std.py", line 1181, in iter
for obj in iterable:
File "E:\ship detection\DRENet-main\utils\datasets.py", line 106, in iter
yield next(self.iterator)
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch\utils\data\dataloader.py", line 630, in next
data = self._next_data()
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch\utils\data\dataloader.py", line 1345, in _next_data
return self._process_data(data)
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch\utils\data\dataloader.py", line 1371, in _process_data
data.reraise()
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch_utils.py", line 694, in reraise
raise exception
RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch\utils\data_utils\worker.py", line 308, in _worker_loop
data = fetcher.fetch(index)
File "C:\Users\17245\anaconda3\envs\pytorch\lib\site-packages\torch\utils\data_utils\fetch.py", line 54, in fetch
return self.collate_fn(data)
File "E:\ship detection\DRENet-main\utils\datasets.py", line 598, in collate_fn
return torch.stack(img, 0), torch.stack(dgimg, 0), torch.cat(label, 0), path, shapes
RuntimeError: stack expects each tensor to be equal size, but got [3, 62, 128] at entry 0 and [3, 67, 128] at entry 1

[WindVChen]

Hi @wenjieo ,

It looks like the error might be due to inconsistent image resolutions in your dataset. One approach is to resize and pad your images to make the resolutions uniform. For instance, if you have images with resolutions 1024x768, 512x256, and 128x256, you could resize the latter two to 1024x512 and 384x768, respectively, and then pad them to reach 1024x768. This will ensure all images have the same final resolution.

Hope this helps.