How to train the original U-Net model with PyTorch?

user_4461300 0
Viewed 1373

I’m trying to implement and train the original U-Net model, but I’m stuck in when I’m trying to train the model using the ISBI Challenge Dataset.

According with the original U-Net model, the network outputs an image with 2 channels and size of 388 x 388. So, my data loader for training generates a tensor with size of [batch, channels=1, width=572, height=572] for the input images and [batch, channels=2, width=388, width=388] for target/output images.

My problem actually is that when I’m trying to use the nn.CrossEntropyLoss() the following error is raised:

RuntimeError: invalid argument 3: only batches of spatial targets supported (3D tensors) but got targets of dimension: 4 at /opt/conda/conda-bld/pytorch_1556653099582/work/aten/src/THNN/generic/SpatialClassNLLCriterion.c:59

I’m just starting with PyTorch (newbie here)… so, I’ll really appreciate if someone could help me to overcome this problem.

The sourcecode is available on GitHub:

https://github.com/dalifreire/cnn_unet_pytorch https://github.com/dalifreire/cnn_unet_pytorch/blob/master/unet_pytorch.ipynb

Best regards!

UPDATE

I just remove the channel dimension from my masks and everything works well… now I’m generating masks with the shape 1 [width=388, height=388].

After that, I’m working with input images (X), target masks (y) and predicted output masks (y_hat) as follow:

X     --> torch.Size([10, 1, 572, 572])
y     --> torch.Size([10, 388, 388])
y_hat --> torch.Size([10, 2, 388, 388])

But, I don’t understand why target masks (y) and predicted masks (y_hat) must have different shapes? It’s so weird for me…

image-segmentation conv-neural-network pytorch python
1 Answers

From the CrossEntropyLoss docstring of PyTorch:

Shape:

    - Input: :math:`(N, C)` where `C = number of classes`, or
      :math:`(N, C, d_1, d_2, ..., d_K)` with :math:`K \geq 1`
      in the case of `K`-dimensional loss.

    - Target: :math:`(N)` where each value is :math:`0 \leq \text{targets}[i] \leq C-1`, or
      :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \geq 1` in the case of
      K-dimensional loss.

    - Output: scalar.
      If :attr:`reduction` is ``'none'``, then the same size as the target:
      :math:`(N)`, or
      :math:`(N, d_1, d_2, ..., d_K)` with :math:`K \geq 1` in the case
      of K-dimensional loss.

If your targets contain the class indices already, you should remove the channel dimension.

Source

user_68471000
Related