Pytorch - 使用 UNet 执行图像去模糊/图像重建答案

【问题标题】：Pytorch - Use a UNet to perform Image Deblurring/Image ReconstructionPytorch - 使用 UNet 执行图像去模糊/图像重建
【发布时间】：2021-08-20 17:51:33
【问题描述】：

目前，我正在处理一个数据集，其中有两种图像：图像的“锐化版本”和相同图像的“模糊版本”，其中综合添加了模糊。我的目标是训练一个模型，该模型将图像的模糊版本纳入并尝试尽可能多地对图像进行去模糊处理，以使“去模糊图像”更接近清晰的版本。在文献中，UNet 架构似乎是一个效果不错的模型。此外，我可以通过 Pytorch (https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/) 使用预训练的 U-Net。

我现在的问题是：当我用我的图像训练这个预训练的 U-Net，然后在我的测试集上尝试它时，我得到以下输出：

原图：

我知道这个预训练模型通常用于生物医学图像分割，但我对如何修改模型以将其用于图像去模糊/重建任务感到很困惑。有人对如何做到这一点有任何建议吗？

如果有任何反馈，我将不胜感激：)

【问题讨论】：

标签： python image image-processing deep-learning pytorch

【解决方案1】：

您使用的 U-net 用于分割（图像的每个像素的分类），而您尝试对图像进行去噪（使图像“更清晰”/去除噪点）。它解释了你得到的结果。

要得到你想要的，正如 DerekG 所说，你首先需要修改输出的通道数。通过修改它，您无法加载所有预训练模型。您必须逐个参数复制参数，直到最后一个。

由于最后一层是随机初始化的，您可以使用您的训练集重新训练模型。您可以冻结或不冻结预训练的部分。

另外，我不确定您的新数据集是什么，但如果它真的与生物医学图像无关，您应该从头开始重新训练您的网络（在这些情况下不应该进行迁移学习），甚至可以更改编码器 -解码器网络。

【讨论】：

【解决方案2】：

来自包含的链接：

import torch
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
    in_channels=3, out_channels=1, init_features=32, pretrained=True)

您定义的模型具有单个输出通道，从而产生灰度图像输出。 RGB 图像需要 3 个输出通道。

【讨论】：

非常感谢您的意见！当我将参数“out_channels”更改为 3 时，不幸的是，我收到以下错误消息：RuntimeError: Error(s) in loading state_dict for UNet: size mismatch for conv.weight: copying a param with shape torch.Size([1, 32, 1, 1]) 从检查点开始，当前模型中的形状是 torch.Size([3, 32, 1, 1])。 conv.bias 的大小不匹配：从检查点复制形状为 torch.Size([1]) 的参数，当前模型中的形状为 torch.Size([3])。
但是关于我的初步结果：我确实将最后一个 Conv2d 层中的 output_channels 从 1 更改为 3，这导致了问题中的结果
是的，这是意料之中的。这是因为您尝试加载的已保存文件假定模型具有与您的情况不同的结构。您可以尝试找到一个不同的预训练模型文件（可能在此处github.com/qubvel/segmentation_models.pytorch），或者您可以修改 state_dict 使其具有与模型相同的参数大小（使用随机权重初始化新参数并在训练期间学习它们）。