Ignore extra channel weights when loading a checkpoint

3DTOPO: <p>I need to load a checkpoint with four channels into an otherwise identical model with three channels and ignore say the 4th channel.</p>
<p>I can’t train the model again since the compute is cost prohibitive. It seems like a fairly common issue with say RGBA &lt;&gt; RGB. It would also be handy for me to know how to use a single channel.</p>
<p>I tried walking over the <code>state_dict</code>:</p>
<pre><code>model.load_state_dict(dict([(n, p) for n, p in checkpoint['state_dict'].items()]), strict=False)
</code></pre>
<p>But still get a mismatch error:</p>
<pre><code>size mismatch for model.vision_model.out.2.weight: copying a param with shape torch.Size([4, 512, 3, 3]) from checkpoint, the shape in current model is torch.Size([3, 512, 3, 3]).
size mismatch for model.vision_model.out.2.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([3]).
</code></pre>
<p>I obviously need to walk each channel but not sure what the best way is. I appreciate any suggestions.</p>
<p>I have very constrained runtime compute and when considering all possible optimizations, this would be the biggest improvement I could make because 25% of inference compute is completely wasted now but having the unneeded weights in GPU memory is my main concern.</p>

ptrblck: <aside class="quote no-group" data-username="3DTOPO" data-post="1" data-topic="168499">
<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/3dtopo/48/27560_2.png" class="avatar"> 3DTOPO:</div>
<blockquote>
<p>I have very constrained runtime compute and when considering all possible optimizations, this would be the biggest improvement I could make because 25% of inference compute is completely wasted now but having the unneeded weights in GPU memory is my main concern.</p>
</blockquote>
</aside>
<p>I’m a bit confused about this statement as I would assume the only difference would be in the inputs having a different number of channels (3 vs. 4).<br>
In this case, only the first conv layer would use a single additional channel, which might be a tiny compute and memory overhead compared to every other layer and operation in this model.<br>
E.g. check this small example:</p>
<pre data-code-wrap="python"><code class="lang-python"># default RGB
modelA = models.resnet50()
x = torch.randn(1, 3, 224, 224)
out = modelA(x) # works

# check number of parameters
nb_paramsA = sum([p.nelement() for p in modelA.parameters()])

# change to RGBA inputs
modelB = models.resnet50()
modelB.conv1 = nn.Conv2d(4, 64, 7, 2, 3, bias=False)
x = torch.randn(1, 4, 224, 224)
out = modelB(x) # works

# check number of parameters
nb_paramsB = sum([p.nelement() for p in modelB.parameters()])

# compare
print((nb_paramsB - nb_paramsA) / nb_paramsB)
# 0.00012269089937124044
</code></pre>
<p>In any case, you could manipulate the <code>state_dict</code> directly by dropping the unwanted channel before calling <code>load_state_dict</code>.</p>

3DTOPO: <p>Wow, thank you for the fascinating example, I’m surprised! I was under the impression that each channel in the model had a set of weights.</p>
<p>IIRC inference is done on a per channel basis though, right? If that is the case then extra channels would still less than ideal.</p>
<p>Do you know of any examples showing how to drop channels from a state_dict?</p>

ptrblck: <aside class="quote no-group" data-username="3DTOPO" data-post="3" data-topic="168499">
<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/3dtopo/48/27560_2.png" class="avatar"> 3DTOPO:</div>
<blockquote>
<p>Wow, thank you for the fascinating example, I’m surprised! I was under the impression that each channel in the model had a set of weights.</p>
</blockquote>
</aside>
<p>The <code>weight</code> attribute of a conv layer has the shape <code>[out_channels, in_channels, height, width]</code>. Since you are defining the <code>out_channels</code> of a conv layer, only the very first one uses the actual input channels of your input tensors.<br>
I.e. look at the first layers of the resnet:</p>
<pre data-code-wrap="python"><code class="lang-python">print(modelA)
ResNet(
  (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (relu): ReLU(inplace=True)
  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
  (layer1): Sequential(
    (0): Bottleneck(
      (conv1): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
      (downsample): Sequential(
        (0): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      )
    )
</code></pre>
<p>As you can see, the output activation of the first conv layer will have 64 channels. The next layers thus expect activations with 64 channels and will define their own <code>out_channels</code> values (e.g. the next two conv layers still outputs 64 channels, while the one afterwards increases it to 256).</p>
<aside class="quote no-group" data-username="3DTOPO" data-post="3" data-topic="168499">
<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/3dtopo/48/27560_2.png" class="avatar"> 3DTOPO:</div>
<blockquote>
<p>IIRC inference is done on a per channel basis though, right? If that is the case then extra channels would still less than ideal.</p>
</blockquote>
</aside>
<p>No, that’s not the case. Each layer will process the entire batch and output the specified activation. The number of channels of your input is only important for the very first conv layer.</p>
<aside class="quote no-group" data-username="3DTOPO" data-post="3" data-topic="168499">
<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/3dtopo/48/27560_2.png" class="avatar"> 3DTOPO:</div>
<blockquote>
<p>Do you know of any examples showing how to drop channels from a state_dict?</p>
</blockquote>
</aside>
<p>Here is one:</p>
<pre data-code-wrap="python"><code class="lang-python"># fails
sd = modelB.state_dict()
modelA.load_state_dict(sd)
# RuntimeError: Error(s) in loading state_dict for ResNet:
#  size mismatch for conv1.weight: copying a param with shape torch.Size([64, 4, 7, 7]) from checkpoint, the shape in current model is torch.Size([64, 3, 7, 7]).

print(sd['conv1.weight'].shape)
# torch.Size([64, 4, 7, 7])

# slice channels
sd['conv1.weight'] = sd['conv1.weight'][:, :3, ...]

# works
modelA.load_state_dict(sd)
# &lt;All keys matched successfully&gt;
</code></pre>