如何用 PyTorch 训练原始 U-Net 模型？

Question

我正在尝试实施和训练 original U-Net model, but I’m stuck in when I’m trying to train the model using the ISBI Challenge Dataset。

根据原始的 U-Net 模型，网络输出一个 2 通道的图像，大小为 388 x 388。因此，我的训练数据加载器生成了一个大小为 [batch,输入图像的通道=1，宽度=572，高度=572]，对于target/output 图片。

我的问题实际上是当我尝试使用 nn.CrossEntropyLoss() 时出现以下错误：

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

我刚开始使用 PyTorch（这里是新手）...所以，如果有人能帮助我解决这个问题，我将不胜感激。

源代码可在 GitHub:

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

此致！

更新

我只是从蒙版中删除了通道尺寸，一切正常......现在我正在生成形状为 1 [width=388，height=388] 的蒙版。

之后，我正在处理输入图像 (X)、目标蒙版 (y) 和预测输出蒙版 (y_hat)，如下所示：

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

但是，我不明白为什么目标掩码（y）和预测掩码（y_hat）必须具有不同的形状？这对我来说太奇怪了……

Answer 1

来自 PyTorch 的 CrossEntropyLoss 文档字符串：

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.

如果您的目标已经包含 class 个索引，您应该删除通道维度。

Source

如何用 PyTorch 训练原始 U-Net 模型？

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

python

image-segmentation

conv-neural-network

pytorch