Get encoder from trained UNet

Flock1: <p>Hi,</p>
<p>I have a trained a UNet model on some images but now, I want to extract the encoder part of the model. My UNet has the following architecture:</p>
<pre><code class="lang-auto">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))
    )
  )
)
</code></pre>
<p>I have tried to load the encoder layers through <code>model.down_convs</code> but I get the following error:</p>
<blockquote>
<p>TypeError                                 Traceback (most recent call last)<br>
 in <br>
----&gt; 1 res = encoder(train_img)</p>
<p>~/anaconda3/envs/work/lib/python3.8/site-packages/torch/nn/modules/module.py in <strong>call</strong>(self, *input, **kwargs)<br>
548             result = self._slow_forward(*input, **kwargs)<br>
549         else:<br>
 → 550             result = self.forward(*input, **kwargs)<br>
551         for hook in self._forward_hooks.values():<br>
552             hook_result = hook(self, input, result)</p>
<p>TypeError: forward() takes 1 positional argument but 2 were given</p>
</blockquote>
<p>What should I do?</p>

suraj.pt: <p>Can you add your code for loading the encoder also?</p>

Flock1: <p>If you are referring to a separate model, then I don’t have any encoder model to load. I have only trained a UNet from which I want to get the encoder weights.</p>

suraj.pt: <p>Are you doing this?</p>
<pre><code class="lang-auto">encoder = model.down_convs
result = encoder(train_img)
</code></pre>
<p>This works for me. It’ll be easier to debug if you post your code so we can reproduce the error</p>

Flock1: <p>Yup, that’s what I am doing. I am loading a <code>.npz</code> file and encoding it.This is the code:</p>
<pre><code class="lang-auto">img = np.load('/HMI20181101_0000_bz.npz', allow_pickle=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_img = img['x']
train_img = train_img.reshape(1, 1, 512, 512)
train_img = np.array(train_img, dtype=np.float)
train_img = torch.from_numpy(train_img)
train_img = train_img.to(device, dtype=torch.float)

res = encoder(train_img)
</code></pre>
<p>This is my model for the encoder:</p>
<pre><code class="lang-auto">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))
  )
)
</code></pre>
<p>And the error I get after running the code is:</p>
<blockquote>
<hr>
<p>TypeError                                 Traceback (most recent call last)<br>
 in <br>
----&gt; 1 res = encoder(train_img)</p>
<p>~/anaconda3/envs/work/lib/python3.8/site-packages/torch/nn/modules/module.py in <strong>call</strong>(self, *input, **kwargs)<br>
548             result = self._slow_forward(*input, **kwargs)<br>
549         else:<br>
 → 550             result = self.forward(*input, **kwargs)<br>
551         for hook in self._forward_hooks.values():<br>
552             hook_result = hook(self, input, result)</p>
<p>TypeError: forward() takes 1 positional argument but 2 were given</p>
</blockquote>
<p>I have attached the <a href="https://gist.github.com/Flock1/34b5eb8685df07bad950868f57577693" rel="noopener nofollow ugc">model</a> so you can try it out. And the weights from <a href="https://github.com/Flock1/solar" rel="noopener nofollow ugc">here</a></p>

Flock1: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, any suggestions?</p>

ptrblck: <p>You cannot call <code>nn.ModuleList</code> directly, as it’s used as a container to store other modules.<br>
If you want to call <code>encoder</code> with an input, you could either try to wrap all modules into <code>nn.Sequential</code> or apply the loop over all modules manually.</p>

Flock1: <p>Are you talking about something like this?</p>
<pre><code class="lang-auto">test_block = torch.nn.Sequential(*list(test_model.children()))[:2]
</code></pre>
<p>Because I get the following model:</p>
<pre><code class="lang-auto">Sequential(
  (0): Conv2d(8, 1, kernel_size=(1, 1), stride=(1, 1))
  (1): 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))
    )
  )
)
</code></pre>
<p>When I run it <code>res = test_block(train_img)</code>, I get the following error:</p>
<blockquote>
<p>RuntimeError: Given groups=1, weight of size [1, 8, 1, 1], expected input[1, 1, 512, 512] to have 8 channels, but got 1 channels instead</p>
</blockquote>

Mughees: <aside class="quote no-group" data-username="Flock1" data-post="8" data-topic="95053">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/flock1/48/12220_2.png" class="avatar"> Flock1:</div>
<blockquote>
<p><code>(0): Conv2d(8, 1, kernel_size=(1, 1), stride=(1, 1))</code></p>
</blockquote>
</aside>
<p>This input layer requires 8 band input. as<br>
<code>Conv2d(in_channels, out_channels, kernel_size)</code> which seems to me the bug so modify the original network to have input layer for 1 band and then extract the encoder.</p>

Flock1: <p>Thanks. I will change the layer. But, I wanted to ask, if is it a problem here, so how come it didn’t create any problems while training since I Have trained on the same type of images?</p>

ptrblck: <p>I would be a bit careful using this code:</p>
<aside class="quote no-group" data-username="Flock1" data-post="8" data-topic="95053">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/flock1/48/12220_2.png" class="avatar"> Flock1:</div>
<blockquote>
<p><code>test_block = torch.nn.Sequential(*list(test_model.children()))[:2]</code></p>
</blockquote>
</aside>
<p>as it will create an <code>nn.Sequential</code> module using the child modules in a sequential order.<br>
In your initial code it seems that</p>
<pre data-code-wrap="python"><code class="lang-python">(conv_final): Conv2d(8, 1, kernel_size=(1, 1), stride=(1, 1))
</code></pre>
<p>is initialized as the first layer, but it supposed to be used as the last layer, which would create this issue.</p>
<p>Pass the <code>nn.ModuleList</code> to <code>nn.Sequential</code> as seen here:</p>
<pre data-code-wrap="python"><code class="lang-python">mlist = nn.ModuleList([
    nn.Linear(1, 2),
    nn.Linear(2, 3),
    nn.Linear(3, 4)])

model = nn.Sequential(*mlist)
model(torch.randn(1, 1))
</code></pre>
<p>or iterate it manually.</p>

Flock1: <p>So, using this,</p>
<aside class="quote no-group" data-username="ptrblck" data-post="11" data-topic="95053">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<pre><code class="lang-auto">mlist = nn.ModuleList([
    nn.Linear(1, 2),
    nn.Linear(2, 3),
    nn.Linear(3, 4)])

model = nn.Sequential(*mlist)
model(torch.randn(1, 1))
</code></pre>
</blockquote>
</aside>
<p>will it be possible to assign UNet encoder parameters to a completely different model? Let’s say if I have weights and biases of the encoder part of the UNet, can I assign those values to a model that I have created using the way you mentioned? Only, in the tuple, I will have the filter size as mentioned.</p>

ptrblck: <p>Yes, that would be possible but you would usually have to load these weights manually or manipulate the <code>state_dict</code> keys so that the parameter names match again.</p>

Flock1: <p>Thanks. I have tried that but I don’t know why, even though I have changed nothing and loading the same model, I get size mismatch error on every layer:</p>
<pre><code class="lang-auto">RuntimeError: Error(s) in loading state_dict for UNet:
	size mismatch for conv_final.weight: copying a param with shape torch.Size([1, 8, 1, 1]) from checkpoint, the shape in current model is torch.Size([1, 64, 1, 1]).
.
.
.
</code></pre>
<p>I don’t know how the number of channels went to 8 from 64.</p>
<p>That’s why I was trying to load the whole model and edit it from there.</p>
<p>For reference, this is the <a href="https://gist.github.com/Flock1/34b5eb8685df07bad950868f57577693" rel="nofollow noopener">model definition</a> and these are the <a href="https://github.com/Flock1/solar" rel="nofollow noopener">weights</a></p>

ptrblck: <p>How was <code>conv_final</code> defined in the original model and how is it defined now?<br>
Could you check the definition of this layer and how you’ve manipulated the <code>state_dict</code>?</p>

Flock1: <p>I guess I was missing an argument value while defining the UNet object, I had to change it from 64 to 8. Like this<br>
<code>test_model = UNet(2, start_filts=8)</code><br>
Now I can try your method to define an encoder.</p>

Flock1: <p>Him</p>
<p>I have been able to extract the encoder and mostly been able to get rid of errors. I followed these steps:</p>
<pre><code class="lang-auto">test_modules = test_model.down_convs
test_modules = nn.Sequential(*test_modules)
test_modules.eval()
print((test_modules._modules))
</code></pre>
<p>And was able to get the model:</p>
<pre><code class="lang-auto">OrderedDict([('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))
))])
</code></pre>
<p>I loaded my image like this (since is <code>.npz</code> file):</p>
<pre><code class="lang-auto">tmp = np.load(file, allow_pickle=True)
img = tmp['x']
img = np.reshape(img,(1,1,512,512))
img = torch.from_numpy(img)
</code></pre>
<p>But I got the error after running <code>test_modules(img)</code>:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
&lt;ipython-input-12-621fab09e884&gt; in &lt;module&gt;
      1 with torch.no_grad():
----&gt; 2     test_modules(img)

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--&gt; 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
    115     def forward(self, input):
    116         for module in self:
--&gt; 117             input = module(input)
    118         return input
    119 

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--&gt; 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

~/Desktop/solar/solar/model.py in forward(self, x)
     65 
     66     def forward(self, x):
---&gt; 67         x = F.relu(self.conv1(x))
     68         x = F.relu(self.conv2(x))
     69         before_pool = x

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    720             result = self._slow_forward(*input, **kwargs)
    721         else:
--&gt; 722             result = self.forward(*input, **kwargs)
    723         for hook in itertools.chain(
    724                 _global_forward_hooks.values(),

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/conv.py in forward(self, input)
    417 
    418     def forward(self, input: Tensor) -&gt; Tensor:
--&gt; 419         return self._conv_forward(input, self.weight)
    420 
    421 class Conv3d(_ConvNd):

~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
    413                             weight, self.bias, self.stride,
    414                             _pair(0), self.dilation, self.groups)
--&gt; 415         return F.conv2d(input, weight, self.bias, self.stride,
    416                         self.padding, self.dilation, self.groups)
    417 

TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not tuple
</code></pre>
<p>The input image to the network is of type <code>&lt;class 'torch.Tensor'&gt;</code>. So what’s happening now?</p>
<p><a class="mention" href="/u/alband">@albanD</a>, you can also let me know.</p>

Flock1: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, <a class="mention" href="/u/alband">@albanD</a>, <a class="mention" href="/u/mughees">@Mughees</a>, <a class="mention" href="/u/suraj.pt">@suraj.pt</a>  please let me know what should I do?</p>

ptrblck: <p>Please don’t tag certain people, as this will demotivate others to answer your question.</p>
<p>Could you post an executable code snippet to reproduce this issue?</p>

Flock1: <aside class="quote no-group" data-username="ptrblck" data-post="19" data-topic="95053">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Please don’t tag certain people, as this will demotivate others to answer your question.</p>
</blockquote>
</aside>
<p>Will keep this in mind from now on.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="19" data-topic="95053">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>Could you post an executable code snippet to reproduce this issue?</p>
</blockquote>
</aside>
<p>I have created a <a href="https://colab.research.google.com/drive/1IvLgNpK1s97lk1waLUEU9jaQpAcEtcFR?usp=sharing" rel="noopener nofollow ugc">Colab</a> notebook, which you can run and reproduce the error.</p>