How to change the output shape of a semantic segmentation model?

username: johnny69
cooked: <p><strong>Context:</strong> the input is a 520x520 image of a cell slide, and the output is a heatmap of where the macrophage cell centres are:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/e/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762_2_690x343.jpeg" alt="image" data-base62-sha1="i5mCGAGf4YHqrRqRjdtLABCF6Cu" width="690" height="343" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762_2_690x343.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/7/e/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/e/7ec261ef39a6ea6ad3ce137b85de6f8e9df59762.jpeg 2x" data-dominant-color="B6A3BD"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">988×492 45.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>This means the input shape is 3x520x520, and the output shape <em>should</em> be 1x520x520.</p>
<p><strong>Problem:</strong> I’m trying to use transfer learning with a pretrained semantic segmentation model to get the output from the input. I’m using <code>deeplabv3_resnet50</code> from <code>torchvision.models.segmentation</code>. However, when I print the model summary with <code>torchinfo.summary()</code>, the input shape is correct (1x3x520x520), but the output shape is something unexpected (either 1x21x520x520 or 1x21x65x65?!)</p>
<pre><code class="lang-auto">==============================================================================================================
Layer (type (var_name))                            Input Shape          Output Shape         Trainable
==============================================================================================================
DeepLabV3 (DeepLabV3)                              [1, 3, 520, 520]     [1, 21, 520, 520]    True
├─IntermediateLayerGetter (backbone)               [1, 3, 520, 520]     [1, 2048, 65, 65]    True
│    └─Conv2d (conv1)                              [1, 3, 520, 520]     [1, 64, 260, 260]    True
│    └─BatchNorm2d (bn1)                           [1, 64, 260, 260]    [1, 64, 260, 260]    True
│    └─ReLU (relu)                                 [1, 64, 260, 260]    [1, 64, 260, 260]    --
│    └─MaxPool2d (maxpool)                         [1, 64, 260, 260]    [1, 64, 130, 130]    --
│    └─Sequential (layer1)                         [1, 64, 130, 130]    [1, 256, 130, 130]   True
│    │    └─Bottleneck (0)                         [1, 64, 130, 130]    [1, 256, 130, 130]   True
│    │    └─Bottleneck (1)                         [1, 256, 130, 130]   [1, 256, 130, 130]   True
│    │    └─Bottleneck (2)                         [1, 256, 130, 130]   [1, 256, 130, 130]   True
│    └─Sequential (layer2)                         [1, 256, 130, 130]   [1, 512, 65, 65]     True
│    │    └─Bottleneck (0)                         [1, 256, 130, 130]   [1, 512, 65, 65]     True
│    │    └─Bottleneck (1)                         [1, 512, 65, 65]     [1, 512, 65, 65]     True
│    │    └─Bottleneck (2)                         [1, 512, 65, 65]     [1, 512, 65, 65]     True
│    │    └─Bottleneck (3)                         [1, 512, 65, 65]     [1, 512, 65, 65]     True
│    └─Sequential (layer3)                         [1, 512, 65, 65]     [1, 1024, 65, 65]    True
│    │    └─Bottleneck (0)                         [1, 512, 65, 65]     [1, 1024, 65, 65]    True
│    │    └─Bottleneck (1)                         [1, 1024, 65, 65]    [1, 1024, 65, 65]    True
│    │    └─Bottleneck (2)                         [1, 1024, 65, 65]    [1, 1024, 65, 65]    True
│    │    └─Bottleneck (3)                         [1, 1024, 65, 65]    [1, 1024, 65, 65]    True
│    │    └─Bottleneck (4)                         [1, 1024, 65, 65]    [1, 1024, 65, 65]    True
│    │    └─Bottleneck (5)                         [1, 1024, 65, 65]    [1, 1024, 65, 65]    True
│    └─Sequential (layer4)                         [1, 1024, 65, 65]    [1, 2048, 65, 65]    True
│    │    └─Bottleneck (0)                         [1, 1024, 65, 65]    [1, 2048, 65, 65]    True
│    │    └─Bottleneck (1)                         [1, 2048, 65, 65]    [1, 2048, 65, 65]    True
│    │    └─Bottleneck (2)                         [1, 2048, 65, 65]    [1, 2048, 65, 65]    True
├─DeepLabHead (classifier)                         [1, 2048, 65, 65]    [1, 21, 65, 65]      True
│    └─ASPP (0)                                    [1, 2048, 65, 65]    [1, 256, 65, 65]     True
│    │    └─ModuleList (convs)                     --                   --                   True
│    │    └─Sequential (project)                   [1, 1280, 65, 65]    [1, 256, 65, 65]     True
│    └─Conv2d (1)                                  [1, 256, 65, 65]     [1, 256, 65, 65]     True
│    └─BatchNorm2d (2)                             [1, 256, 65, 65]     [1, 256, 65, 65]     True
│    └─ReLU (3)                                    [1, 256, 65, 65]     [1, 256, 65, 65]     --
│    └─Conv2d (4)                                  [1, 256, 65, 65]     [1, 21, 65, 65]      True
├─FCNHead (aux_classifier)                         [1, 1024, 65, 65]    [1, 21, 65, 65]      True
│    └─Conv2d (0)                                  [1, 1024, 65, 65]    [1, 256, 65, 65]     True
│    └─BatchNorm2d (1)                             [1, 256, 65, 65]     [1, 256, 65, 65]     True
│    └─ReLU (2)                                    [1, 256, 65, 65]     [1, 256, 65, 65]     --
│    └─Dropout (3)                                 [1, 256, 65, 65]     [1, 256, 65, 65]     --
│    └─Conv2d (4)                                  [1, 256, 65, 65]     [1, 21, 65, 65]      True
==============================================================================================================
Total params: 42,004,074
Trainable params: 42,004,074
Non-trainable params: 0
Total mult-adds (G): 178.72
==============================================================================================================
Input size (MB): 3.24
Forward/backward pass size (MB): 2294.42
Params size (MB): 168.02
Estimated Total Size (MB): 2465.68
==============================================================================================================
</code></pre>
<p>Does anyone know why the output shape is turning out like that and how to fix it so that it’s 1x1x520x520?</p>

username: ptrblck
cooked: <p>As explained in <a href="https://arxiv.org/pdf/1706.05587.pdf">the paper</a> the spatial size of the output depends on the <code>output_stride</code>:</p>
<blockquote>
<p>Recall that output stride is defined as the ratio of input image spatial resolution to final output resolution. For example, when output stride = 8, the last two blocks (block3 and block4 in our notation) in the original ResNet contains atrous convolution with rate = 2 and rate = 4 respectively.</p>
</blockquote>
<p>And it seems the logits were upsampled:</p>
<blockquote>
<p>Upsampling logits: In our previous works [ 10 , 11 ], the target groundtruths are downsampled by 8 during training when output stride = 8. We find it important to keep the groundtruths intact and instead upsample the final logits, since downsampling the groundtruths removes the fine annotations resulting in no back-propagation of details.</p>
</blockquote>
<p>The 21 output channels should represent 20 classes + (I guess) 1 background class as given in the <a href="https://pytorch.org/vision/main/models/generated/torchvision.models.segmentation.deeplabv3_resnet50.html#torchvision.models.segmentation.deeplabv3_resnet50">docs</a>:</p>
<blockquote>
<p>These weights were trained on a subset of COCO, using only the 20 categories that are present in the Pascal VOC dataset. Also available as <code>DeepLabV3_ResNet50_Weights.DEFAULT</code>.</p>
</blockquote>

username: johnny69
cooked: <p>Okay, starting to make more sense but I’m still a bit confused - so what should I do if I want the output image to have a shape of 520x520?</p>
<p>And how do I change the number of classes (I tried setting the <code>num_classes</code> parameter in the model to <code>1</code> but that gets an error saying that the model expected <code>21</code> for that parameter).</p>

username: mMagmer
cooked: <p>You are trying to load pre-trained model weight to it.<br>
Don’t need pre-trained weight<br>
use :<br>
<code>deeplabv3_resnet50(weights=None,num_classes=1)</code><br>
and if it’s necessary, try:</p><aside class="quote" data-post="2" data-topic="51004">
  <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">
    <a href="https://discuss.pytorch.org/t/removing-classification-layer-for-resnet101-deeplabv3/51004/2">Removing classification layer for resnet101-deeplabv3</a> 
  </div>
  <blockquote>
    Would you like to remove the very last layer in this model or the complete classification head (DeepLabHead)? 
In the former case, you could just set model.classifier[4] to an nn.Identity layer: 
model = models.segmentation.deeplabv3_resnet101()
model.classifier[4] = nn.Identity()

If you would like to change this layer for a custom one using a new number of classes, just replace it with 
nn.Conv2d(
    in_channels=256,
    out_channels=nb_classes,
    kernel_size=1,
    stride=1
)
  </blockquote>
</aside>


username: johnny69
cooked: <p>Thanks, I just ended up just removing the last layer and it worked great.</p>