Change the dimensions of the UNET network output

username: Firas_moh
cooked: <p>As I know, the UNET network predicts the outputs with the same size as the input. How can I change the output of the UNET to be in different sizes (I want certain dimensions for the height and the width to be compatible with my ground truth images)?<br>
please, how can I do this? I want the change to be inside the network? not after the prediction? since when I did the resizing on the prediction image when I load the model again for the testing process, it give a mismatching error.</p>

username: ksz
cooked: <p>Your description is a little chaotic for me, sorry. Please try to describe your problem better.</p>
<p>Two ideas:</p>
<ul>
<li>you can resize the output from the U-NET to yours new dimensions - independently from the PyTorch</li>
<li>you can modify last layers form the U-NET to return specific dimensions. However, to help you with this we need code you work with and furthers explanations.</li>
</ul>

username: Firas_moh
cooked: <p>Thank you so much for your help. I am sorry for that I am still new in this area.</p>
<p>I changed the last layer for Unet to be 1/5 from the input image as in the following code:</p>
<pre><code class="lang-auto">
class BaseConv(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, padding,

                 stride):

        super(BaseConv, self).__init__()

        self.act = nn.ReLU()

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size, padding,

                               stride)

        

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size,

                               padding, stride)

    def forward(self, x):

        x = self.act(self.conv1(x))

        x = self.act(self.conv2(x))

        return x

class DownConv(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, padding,

                 stride):

        super(DownConv, self).__init__()

        self.pool1 = nn.MaxPool2d(kernel_size=2)

        self.conv_block = BaseConv(in_channels, out_channels, kernel_size,

                                   padding, stride)

    def forward(self, x):

        x = self.pool1(x)

        x = self.conv_block(x)

        return x

class UpConv(nn.Module):

    def __init__(self, in_channels, in_channels_skip, out_channels,

                 kernel_size, padding, stride):

        super(UpConv, self).__init__()

        self.conv_trans1 = nn.ConvTranspose2d(

            in_channels, in_channels, kernel_size=2, padding=0, stride=2)

        self.conv_block = BaseConv(

            in_channels=in_channels + in_channels_skip,

            out_channels=out_channels,

            kernel_size=kernel_size,

            padding=padding,

            stride=stride)

    def forward(self, x, x_skip):

        x = self.conv_trans1(x)

        x = torch.cat((x, x_skip), dim=1)

        x = self.conv_block(x)

        return x

class UNet(nn.Module):

    def __init__(self, in_channels, out_channels, n_class, kernel_size,

                 padding, stride):

        super(UNet, self).__init__()

        self.init_conv = BaseConv(in_channels, out_channels, kernel_size,

                                  padding, stride)

        self.down1 = DownConv(out_channels, 2 * out_channels, kernel_size,

                              padding, stride)

        self.down2 = DownConv(2 * out_channels, 4 * out_channels, kernel_size,

                              padding, stride)

        self.down3 = DownConv(4 * out_channels, 8 * out_channels, kernel_size,

                              padding, stride)

        self.up3 = UpConv(8 * out_channels, 4 * out_channels, 4 * out_channels,

                          kernel_size, padding, stride)

        self.up2 = UpConv(4 * out_channels, 2 * out_channels, 2 * out_channels,

                          kernel_size, padding, stride)

        self.up1 = UpConv(2 * out_channels, out_channels, out_channels,

                          kernel_size, padding, stride)

        self.out = nn.Conv2d(out_channels, n_class, kernel_size, padding, stride)

    def forward(self, x):

        # Encoder

        x = self.init_conv(x)

        x1 = self.down1(x)

        x2 = self.down2(x1)

        x3 = self.down3(x2)

        # Decoder

        x_up = self.up3(x3, x2)

        x_up = self.up2(x_up, x1)

        x_up = self.up1(x_up, x)

        x_up = nn.functional.interpolate(x_up, scale_factor=0.5, mode='bilinear', align_corners=True)

        # x_out = F.log_softmax(self.out(x_up), 1)

        return x_up
</code></pre>
<p>and I trained the UNET model using images with size [1,3,256,256], I don’t change any thing just save the model. After the training process, I load the model as in the following:</p>
<pre><code class="lang-auto">model = UNet(in_channels=3,
                      out_channels=1,
                      n_class=1,
                      kernel_size=3,
                      padding=1,
                      stride=1)

model.load_state_dict(torch.load(model_param_path), strict=False)
</code></pre>
<p>but I got the following error:</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-5-f070578da7ca&gt; in &lt;module&gt;
      1 model = model_name.to(device)
      2 
----&gt; 3 model.load_state_dict(torch.load(model_param_path), strict=False)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1043         if len(error_msgs) &gt; 0:
   1044             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-&gt; 1045                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1046         return _IncompatibleKeys(missing_keys, unexpected_keys)
   1047 

RuntimeError: Error(s) in loading state_dict for UNet:
	size mismatch for init_conv.conv1.weight: copying a param with shape torch.Size([1, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 3, 3, 3]).
	size mismatch for init_conv.conv1.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([64]).
	size mismatch for init_conv.conv2.weight: copying a param with shape torch.Size([1, 1, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 3]).
	size mismatch for init_conv.conv2.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([64]).
	size mismatch for down1.conv_block.conv1.weight: copying a param with shape torch.Size([2, 1, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 64, 3, 3]).
	size mismatch for down1.conv_block.conv1.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for down1.conv_block.conv2.weight: copying a param with shape torch.Size([2, 2, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 128, 3, 3]).
	size mismatch for down1.conv_block.conv2.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for down2.conv_block.conv1.weight: copying a param with shape torch.Size([4, 2, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 128, 3, 3]).
	size mismatch for down2.conv_block.conv1.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([256]).
	size mismatch for down2.conv_block.conv2.weight: copying a param with shape torch.Size([4, 4, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 256, 3, 3]).
	size mismatch for down2.conv_block.conv2.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([256]).
	size mismatch for down3.conv_block.conv1.weight: copying a param with shape torch.Size([8, 4, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 256, 3, 3]).
	size mismatch for down3.conv_block.conv1.bias: copying a param with shape torch.Size([8]) from checkpoint, the shape in current model is torch.Size([512]).
	size mismatch for down3.conv_block.conv2.weight: copying a param with shape torch.Size([8, 8, 3, 3]) from checkpoint, the shape in current model is torch.Size([512, 512, 3, 3]).
	size mismatch for down3.conv_block.conv2.bias: copying a param with shape torch.Size([8]) from checkpoint, the shape in current model is torch.Size([512]).
	size mismatch for up3.conv_trans1.weight: copying a param with shape torch.Size([8, 8, 2, 2]) from checkpoint, the shape in current model is torch.Size([512, 512, 2, 2]).
	size mismatch for up3.conv_trans1.bias: copying a param with shape torch.Size([8]) from checkpoint, the shape in current model is torch.Size([512]).
	size mismatch for up3.conv_block.conv1.weight: copying a param with shape torch.Size([4, 12, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 768, 3, 3]).
	size mismatch for up3.conv_block.conv1.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([256]).
	size mismatch for up3.conv_block.conv2.weight: copying a param with shape torch.Size([4, 4, 3, 3]) from checkpoint, the shape in current model is torch.Size([256, 256, 3, 3]).
	size mismatch for up3.conv_block.conv2.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([256]).
	size mismatch for up2.conv_trans1.weight: copying a param with shape torch.Size([4, 4, 2, 2]) from checkpoint, the shape in current model is torch.Size([256, 256, 2, 2]).
	size mismatch for up2.conv_trans1.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([256]).
	size mismatch for up2.conv_block.conv1.weight: copying a param with shape torch.Size([2, 6, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 384, 3, 3]).
	size mismatch for up2.conv_block.conv1.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for up2.conv_block.conv2.weight: copying a param with shape torch.Size([2, 2, 3, 3]) from checkpoint, the shape in current model is torch.Size([128, 128, 3, 3]).
	size mismatch for up2.conv_block.conv2.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for up1.conv_trans1.weight: copying a param with shape torch.Size([2, 2, 2, 2]) from checkpoint, the shape in current model is torch.Size([128, 128, 2, 2]).
	size mismatch for up1.conv_trans1.bias: copying a param with shape torch.Size([2]) from checkpoint, the shape in current model is torch.Size([128]).
	size mismatch for up1.conv_block.conv1.weight: copying a param with shape torch.Size([1, 3, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 192, 3, 3]).
	size mismatch for up1.conv_block.conv1.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([64]).
	size mismatch for up1.conv_block.conv2.weight: copying a param with shape torch.Size([1, 1, 3, 3]) from checkpoint, the shape in current model is torch.Size([64, 64, 3, 3]).
	size mismatch for up1.conv_block.conv2.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([64]).
	size mismatch for out.weight: copying a param with shape torch.Size([1, 1, 3, 3]) from checkpoint, the shape in current model is torch.Size([1, 64, 3, 3]).

</code></pre>

username: ksz
cooked: <p>Please give us fully-reproducible code - the minimal example with saving and loading the model.</p>
<p>Here is working example:</p>
<pre><code class="lang-auto">model = UNet(in_channels=3,
                      out_channels=1,
                      n_class=1,
                      kernel_size=3,
                      padding=1,
                      stride=1)
torch.save(model.state_dict(), "./test.pt")


new_model = model = UNet(in_channels=3,
                      out_channels=1,
                      n_class=1,
                      kernel_size=3,
                      padding=1,
                      stride=1)
new_model.load_state_dict(torch.load("./test.pt"))

x = torch.rand((1, 3, 224, 224))
print(model(x).equal(new_model(x)))
</code></pre>

username: Firas_moh
cooked: <p>Thank you so much for your help. I solved the problem. The problem was in saving the model. My model works fine now with using</p>
<pre><code class="lang-auto">torch.save(model.state_dict(), PATH)
</code></pre>
<p>but when I use a general checkpoint for inference using</p>
<pre><code class="lang-auto">torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)
</code></pre>
<p>the error of size mismatching and missing keys occur.</p>