从训练有素的UNet获取编码器

问题描述

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

count("*")

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

TypeError Traceback（最近一次通话最近） ----> 1个res =编码器（train_img）

〜/ anaconda3 / envs / work / lib / python3.8 / site-packages / torch / nn / modules / module.py 在通话中（self，* input，** kwargs）548 result = self._slow_forward（* input，** kwargs）549 else：–> 550结果= self.forward（* input，** kwargs）551挂钩 self._forward_hooks.values（）：552 hook_result = hook（自我，输入， 结果）

TypeError：forward（）接受1个位置参数，但给出了2个

我已附上model，因此您可以尝试一下。还有here

中的权重

请让我知道。

解决方法

这应该有效。

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.