ROI size in DynNet during inference

[talhaanwarch]

DynNet model copied from

MONAI/tests/test_dynunet.py

Line 62 in 7f42efe

"spatial_dims": 3,

DynUNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=5,
        kernel_size=(3, (1, 1, 3), 3, 3),
        strides=((1, 2, 1), 2, 2, 1),
        upsample_kernel_size=(2, 2, 1),
        res_block=False,
    )

def validation_step(self,batch,batch_idx)
  image,segment=batch["image"], batch["label"]
  print(image.shape,segment.shape)
  roi_size = (96, 96, 96)
  sw_batch_size = 4
  outputs = sliding_window_inference(image, roi_size, sw_batch_size, self.forward)

Shape printed
torch.Size([1, 1, 358, 353, 201]) torch.Size([1, 1, 358, 353, 201])
But I got this error when passed to sliding_window_inference
RuntimeError: The size of tensor a (96) must match the size of tensor b (48) at non-singleton dimension 2

[Nic-Ma]

Hi @talhaanwarch ,

Could you please help share the full error log for debug?
I think maybe the issue was not from sliding_window_inference?

Thanks.

[talhaanwarch]

please check full error logs

RuntimeError                              Traceback (most recent call last)
<ipython-input-48-7114b32222fa> in <module>()
----> 1 get_ipython().run_cell_magic('time', '', 'trainer.fit(model)')

19 frames
/usr/local/lib/python3.7/dist-packages/IPython/core/interactiveshell.py in run_cell_magic(self, magic_name, line, cell)
   2115             magic_arg_s = self.var_expand(line, stack_depth)
   2116             with self.builtin_trap:
-> 2117                 result = fn(magic_arg_s, cell)
   2118             return result
   2119 

<decorator-gen-53> in time(self, line, cell, local_ns)

/usr/local/lib/python3.7/dist-packages/IPython/core/magic.py in <lambda>(f, *a, **k)
    186     # but it's overkill for just that one bit of state.
    187     def magic_deco(arg):
--> 188         call = lambda f, *a, **k: f(*a, **k)
    189 
    190         if callable(arg):

/usr/local/lib/python3.7/dist-packages/IPython/core/magics/execution.py in time(self, line, cell, local_ns)
   1187         if mode=='eval':
   1188             st = clock2()
-> 1189             out = eval(code, glob, local_ns)
   1190             end = clock2()
   1191         else:

<timed eval> in <module>()

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloaders, val_dataloaders, datamodule, train_dataloader)
    551         self.checkpoint_connector.resume_start()
    552 
--> 553         self._run(model)
    554 
    555         assert self.state.stopped

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model)
    916 
    917         # dispatch `start_training` or `start_evaluating` or `start_predicting`
--> 918         self._dispatch()
    919 
    920         # plugin will finalized fitting (e.g. ddp_spawn will load trained model)

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _dispatch(self)
    984             self.accelerator.start_predicting(self)
    985         else:
--> 986             self.accelerator.start_training(self)
    987 
    988     def run_stage(self):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/accelerator.py in start_training(self, trainer)
     90 
     91     def start_training(self, trainer: "pl.Trainer") -> None:
---> 92         self.training_type_plugin.start_training(trainer)
     93 
     94     def start_evaluating(self, trainer: "pl.Trainer") -> None:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in start_training(self, trainer)
    159     def start_training(self, trainer: "pl.Trainer") -> None:
    160         # double dispatch to initiate the training loop
--> 161         self._results = trainer.run_stage()
    162 
    163     def start_evaluating(self, trainer: "pl.Trainer") -> None:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in run_stage(self)
    994         if self.predicting:
    995             return self._run_predict()
--> 996         return self._run_train()
    997 
    998     def _pre_training_routine(self):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run_train(self)
   1029             self.progress_bar_callback.disable()
   1030 
-> 1031         self._run_sanity_check(self.lightning_module)
   1032 
   1033         # enable train mode

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run_sanity_check(self, ref_model)
   1113             # run eval step
   1114             with torch.no_grad():
-> 1115                 self._evaluation_loop.run()
   1116 
   1117             self.on_sanity_check_end()

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py in run(self, *args, **kwargs)
    109             try:
    110                 self.on_advance_start(*args, **kwargs)
--> 111                 self.advance(*args, **kwargs)
    112                 self.on_advance_end()
    113                 self.iteration_count += 1

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py in advance(self, *args, **kwargs)
    109 
    110         dl_outputs = self.epoch_loop.run(
--> 111             dataloader_iter, self.current_dataloader_idx, dl_max_batches, self.num_dataloaders
    112         )
    113 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/base.py in run(self, *args, **kwargs)
    109             try:
    110                 self.on_advance_start(*args, **kwargs)
--> 111                 self.advance(*args, **kwargs)
    112                 self.on_advance_end()
    113                 self.iteration_count += 1

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py in advance(self, dataloader_iter, dataloader_idx, dl_max_batches, num_dataloaders)
    108         # lightning module methods
    109         with self.trainer.profiler.profile("evaluation_step_and_end"):
--> 110             output = self.evaluation_step(batch, batch_idx, dataloader_idx)
    111             output = self.evaluation_step_end(output)
    112 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py in evaluation_step(self, batch, batch_idx, dataloader_idx)
    152             self.trainer.lightning_module._current_fx_name = "validation_step"
    153             with self.trainer.profiler.profile("validation_step"):
--> 154                 output = self.trainer.accelerator.validation_step(step_kwargs)
    155 
    156         return output

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/accelerators/accelerator.py in validation_step(self, step_kwargs)
    209         """
    210         with self.precision_plugin.val_step_context(), self.training_type_plugin.val_step_context():
--> 211             return self.training_type_plugin.validation_step(*step_kwargs.values())
    212 
    213     def test_step(self, step_kwargs: Dict[str, Union[Any, int]]) -> Optional[STEP_OUTPUT]:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py in validation_step(self, *args, **kwargs)
    176 
    177     def validation_step(self, *args, **kwargs):
--> 178         return self.model.validation_step(*args, **kwargs)
    179 
    180     def test_step(self, *args, **kwargs):

<ipython-input-41-575a84b9f0c5> in validation_step(self, batch, batch_idx)
     68     roi_size = (160, 160, 160)
     69     sw_batch_size = 4
---> 70     outputs = sliding_window_inference(image, roi_size, sw_batch_size, self.forward)
     71     loss=self.criterion(outputs,segment)
     72     outputs = [self.post_pred(i) for i in decollate_batch(outputs)]

/usr/local/lib/python3.7/dist-packages/monai/inferers/utils.py in sliding_window_inference(inputs, roi_size, sw_batch_size, predictor, overlap, mode, sigma_scale, padding_mode, cval, sw_device, device, *args, **kwargs)
    140         # store the result in the proper location of the full output. Apply weights from importance map.
    141         for idx, original_idx in zip(slice_range, unravel_slice):
--> 142             output_image[original_idx] += importance_map * seg_prob[idx - slice_g]
    143             count_map[original_idx] += importance_map
    144 

RuntimeError: The size of tensor a (160) must match the size of tensor b (80) at non-singleton dimension 2

[Nic-Ma]

Hi @talhaanwarch ,

I didn't see any obvious issue in your log..Could you please help share your full program?
Maybe @yiheng-wang-nv and I can reproduce the issue locally with it.

Thanks in advance.

[talhaanwarch]

@Nic-Ma i reproduced this error in speen data using the tutorial given.
https://colab.research.google.com/gist/talhaanwarch/0a91d1b506e901db3feacaad45a07ea7/spleen_segmentation_3d_lightning.ipynb

[Nic-Ma]

Hi @talhaanwarch ,

I think the root cause is that your network args(strides) make the output data to have a different shape as input data.
I will try to enhance MONAI sliding window to support arbitrary output shape later, but for your current issue, you can change the network args to:

self._model = DynUNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=2,
        kernel_size=(3, (1, 1, 3), 3, 3),
        strides=(1, 2, 2, 1),
        upsample_kernel_size=(2, 2, 1),
        res_block=False,
    )

Thanks.