tensor size mismatch error

[AkankshaP0102]

I am trying to train STEGO on a custom dataset but during the training process if I provide labels for the corresponding images I get the following error: Traceback (most recent call last): File "train_segmentation.py", line 598, in my_app trainer.fit(model, train_loader, val_loader) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 741, in fit self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt return trainer_fn(*args, kwargs) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl self._run(model, ckpt_path=ckpt_path) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1199, in _run self._dispatch() File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in _dispatch self.training_type_plugin.start_training(self) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training self._results = trainer.run_stage() File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage return self._run_train() File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1311, in _run_train self._run_sanity_check(self.lightning_module) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1375, in _run_sanity_check self._evaluation_loop.run() File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run self.advance(*args, *kwargs) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 110, in advance dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx, dl_max_batches, self.num_dataloaders) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 145, in run self.advance(args, kwargs) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 122, in advance output = self._evaluation_step(batch, batch_idx, dataloader_idx) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py", line 217, in _evaluation_step output = self.trainer.accelerator.validation_step(step_kwargs) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/accelerators/accelerator.py", line 239, in validation_step return self.training_type_plugin.validation_step(step_kwargs.values()) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 219, in validation_step return self.model.validation_step(args, *kwargs) File "train_segmentation.py", line 354, in validation_step self.linear_metrics.update(linear_preds, label) File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/torchmetrics/metric.py", line 405, in wrapped_func raise err File "/media/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/torchmetrics/metric.py", line 395, in wrapped_func update(args, **kwargs) File "/media/2d46715b-293d-4478-acd4-5f000d443896/stego-studies/src/utils.py", line 240, in update mask = (actual >= 0) & (actual < self.n_classes) & (preds >= 0) & (preds < self.n_classes) RuntimeError: The size of tensor a (3276800) must match the size of tensor b (10240) at non-singleton dimension 0

Please help me with the following. Thank you

[Ruhrozz]

I have found the solution

Most likely you have the same issue as mine. In DirectoryDataset after Image.open("mask") you have mask shape [H, W, 3]. But after that in validation step it interpolates to labels[-2:], so now predict shape is something like [B, C, H, 3].

I have solved the problem by converting mask to grayscale so that mask shape is just [H, W]:

label = Image.open(join(self.label_dir, label_fn)).convert('L')

See this and this for additional information.

[AkankshaP0102]

Thank you @Ruhrozz .I tried the above solution suggested by you. It has solved the tensor error but I'm facing issue for the lable image shape. Was this type of error faced by you?? Your help will be appreciated. I have pasted the error below:

Traceback (most recent call last): File "train_segmentation.py", line 503, in my_app trainer.fit(model, train_loader, val_loader) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 741, in fit self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 685, in _call_and_handle_interrupt return trainer_fn(*args, kwargs) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 777, in _fit_impl self._run(model, ckpt_path=ckpt_path) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1199, in _run self._dispatch() File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1279, in _dispatch self.training_type_plugin.start_training(self) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 202, in start_training self._results = trainer.run_stage() File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1289, in run_stage return self._run_train() File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1311, in _run_train self._run_sanity_check(self.lightning_module) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1375, in _run_sanity_check self._evaluation_loop.run() File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/base.py", line 151, in run output = self.on_run_end() File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 131, in on_run_end self._evaluation_epoch_end(outputs) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py", line 236, in _evaluation_epoch_end model.validation_epoch_end(outputs) File "train_segmentation.py", line 297, in validation_epoch_end ax[1, i].imshow(self.label_cmap[output["label"][i]]) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/matplotlib/_api/deprecation.py", line 459, in wrapper return func(*args, *kwargs) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/matplotlib/init.py", line 1414, in inner return func(ax, map(sanitize_sequence, args), kwargs) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/matplotlib/axes/_axes.py", line 5487, in imshow im.set_data(X) File "/media/emsg/2d46715b-293d-4478-acd4-5f000d443896/anaconda3/envs/stegostudies/lib/python3.7/site-packages/matplotlib/image.py", line 716, in set_data .format(self._A.shape)) TypeError: Invalid shape (1, 320, 320, 3) for image data