0

我已经在一些图像上训练了一个 UNet 模型，但现在，我想提取模型的编码器部分。我的 UNet 具有以下架构：

UNet(
  (conv_final): Conv2d(8, 1, kernel_size=(1, 1), stride=(1, 1))
  (down_convs): ModuleList(
    (0): DownConv(
      (conv1): Conv2d(1, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (1): DownConv(
      (conv1): Conv2d(8, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (2): DownConv(
      (conv1): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (3): DownConv(
      (conv1): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    )
    (4): DownConv(
      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
  )
  (up_convs): ModuleList(
    (0): UpConv(
      (upconv): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))
      (conv1): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
    (1): UpConv(
      (upconv): ConvTranspose2d(64, 32, kernel_size=(2, 2), stride=(2, 2))
      (conv1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
    (2): UpConv(
      (upconv): ConvTranspose2d(32, 16, kernel_size=(2, 2), stride=(2, 2))
      (conv1): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
    (3): UpConv(
      (upconv): ConvTranspose2d(16, 8, kernel_size=(2, 2), stride=(2, 2))
      (conv1): Conv2d(16, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
      (conv2): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    )
  )
)

我试图通过 model.down_convs 加载编码器层，但出现以下错误：

----> 1 res = encoder(train_img) 中的 TypeError Traceback（最近一次调用最后一次）

~/anaconda3/envs/work/lib/python3.8/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs) 548 结果 = self._slow_forward(*input, * *kwargs）549 其他：-> 550 结果 = self.forward(*input, **kwargs) 551 for hook in self._forward_hooks.values(): 552 hook_result = hook(self, input, result)

TypeError: forward() 接受 1 个位置参数，但给出了 2 个

我已附上模型，您可以尝试一下。还有这里的权重

请告诉我。

4

1 回答 1

0

这应该有效。

net = UNet(8) # network object having 8 classes
net.load_state_dict(torch.load('PATH'))
print(net) #see the names of the layers of encoder. 
net1 = net.down_convs #as you have named the encoder as down_convs

#net1 is your encoder.
于 2020-09-07T07:07:10.253 回答