How to convert a pre-trained model from NCHW to NHWC format

username: _singh
cooked: <p>Using a pre-trained PyTorch model (N <strong>C</strong> HW format) but my acceleration platform requires model in NHW<strong>C</strong> format.</p>
<p>Is there an easy way for converting PyTorch model to NHWC format?</p>
<p>I have permuted the weights by fetching weights from PyTorch’s <code>state_dict()</code> method like -</p>
<pre><code class="lang-auto">
if('conv' in str(key)):
            params[key] = value.permute(0,2,3,1)
</code></pre>
<p>But am unable to repopulate the model with permuted dictionary - <code>model.load_state_dict(Updated_params)</code> gives -</p>
<pre><code class="lang-auto">size mismatch for stage4.2.branches.2.3.conv1.weight: copying a param with shape torch.Size([128, 3,
3, 128]) from checkpoint, the shape in current model is torch.Size([128, 128, 3, 3]).
</code></pre>
<p>To resolve this, how can I define layers in the new model in NHWC format in PyTorch.</p>
<p>The link below doesn’t seem to solve this issue.<br>
<a href="https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html#beta-channels-last-memory-format-in-pytorch" rel="noopener nofollow ugc">Channels Last Memory Format in PyTorch</a></p>
<p>Thanks!</p>

username: Nikronic
cooked: <p>Hi,<br>
You are using wrong permutation order.</p>
<aside class="quote no-group" data-username="_singh" data-post="1" data-topic="97589">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c37758/48.png" class="avatar"> _singh:</div>
<blockquote>
<p><code>value.permute(0,2,3,1)</code></p>
</blockquote>
</aside>
<p>See this example:</p>
<pre><code class="lang-auto">x = torch.randn(1, 128, 128, 3)
# your order
x.permute(0,2,3,1).shape # torch.Size([1, 128, 3, 128])

# correct order
x = x.permute(0, 3, 1, 2)
x.shape  # torch.Size([1, 3, 128, 128])
</code></pre>
<p>And the error corresponds to this issue.</p>
<p>I am not sure still it could work fine or not because of channel mutation, because of the forward method. For instance, concatenation, squeezing, etc and all other method which use dim argument to do the operation, if exist in the forward function, may cause issues. You may need to override forward function w.r.t. channel changes.</p>
<p>Bests</p>

username: _singh
cooked: <p>Hi,<br>
Thanks for your reply. If you see, you are recommending the solution other way around i.e. NHWC to NCHW. I want the opposite.</p>
<p>I want (1, 3,128,128) to (1, 128, 128, 3) for which</p>
<blockquote>
<p>value.permute(0,2,3,1)</p>
</blockquote>
<p>is the correct order.</p>
<p>Thanks for the insight regarding forward function though.</p>

username: Nikronic
cooked: <aside class="quote no-group" data-username="_singh" data-post="3" data-topic="97589">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c37758/48.png" class="avatar"> _singh:</div>
<blockquote>
<p>I want the opposite</p>
</blockquote>
</aside>
<p>Ow sorry! I misunderstood the weight with input.</p>
<aside class="quote no-group" data-username="_singh" data-post="1" data-topic="97589">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/_/c37758/48.png" class="avatar"> _singh:</div>
<blockquote>
<p>The link below doesn’t seem to solve this issue.<br>
<a href="https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html#beta-channels-last-memory-format-in-pytorch" rel="noopener nofollow ugc">Channels Last Memory Format in PyTorch</a></p>
</blockquote>
</aside>
<p>What was the issue with this tutorial? I think it is ok as you can send all layers to channel last or channel first mode. I am not sure I am missing something here.</p>
<pre><code class="lang-auto"># define channel first
conv = nn.Conv2d(128, 128, 3)  # replicate [128, 128, 3, 3] weight tensor
print(conv.weight.shape)
print(conv.weight.stride())

# convert to channel last
conv = conv.to(memory_format=torch.channels_last)
print(conv.weight.shape)
print(conv.weight.stride())

# convert back to channel first
conv = conv.to(memory_format=torch.contiguous_format)
print(conv.weight.shape)
print(conv.weight.stride())

# output
# torch.Size([128, 128, 3, 3])
# (1152, 9, 3, 1)
# torch.Size([128, 128, 3, 3])
# (1152, 1, 384, 128)
# torch.Size([128, 128, 3, 3])
# (1152, 9, 3, 1)
</code></pre>

username: VitalyFedyunin
cooked: <aside class="quote no-group" data-username="Nikronic" data-post="4" data-topic="97589">
<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/nikronic/48/71873_2.png" class="avatar"> Nikronic:</div>
<blockquote>
<pre><code class="lang-auto">conv = conv.to(memory_format=torch.channels_last)
print(conv.weight.shape)
print(conv.weight.stride())
</code></pre>
</blockquote>
</aside>
<p>This is the correct way to convert the existing model or layer. Please also make sure you are converting inputs as well</p>
<pre><code class="lang-auto">input = input.to(memory_format=torch.channels_last)
</code></pre>

username: _singh
cooked: <p>Please see below code snippet and comments in print lines:</p>
<pre><code class="lang-auto">device = 'cuda'
input = torch.randint(1, 10, (2, 8, 4, 4), dtype=torch.float32, device= device, requires_grad=True)
model = torch.nn.Conv2d(8, 4, 3)

print(input.shape)  # torch.Size([2, 8, 4, 4])
input = input.contiguous(memory_format=torch.channels_last)
print(input.shape)  # still torch.Size([2, 8, 4, 4]) - need [2, 4, 4, 8] (NHWC) here
model = model.to(memory_format=torch.channels_last)
print(model)    # output= Conv2d(8, 4, kernel_size=(3, 3), stride=(1, 1))
model = model.to(device)
out = model(input)
print(out.shape)    # output= torch.Size([2, 4, 2, 2])  | need [2, 2, 2, 4] (NHWC) here 
print(out.is_contiguous(memory_format=torch.channels_last)) # Output: True
</code></pre>

username: _singh
cooked: <p>Please see <a href="https://discuss.pytorch.org/t/how-to-convert-a-pre-trained-model-from-nchw-to-nhwc-format/97589/6">this</a> comment.</p>
<p>It appears <code>memory_format = torch.channels_last</code>  is not converting the layers/input to NHWC format. It performs a different functionality. <a href="https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html#beta-channels-last-memory-format-in-pytorch" rel="noopener nofollow ugc">Pytorch Channels Last Memory Format</a> webpage also doesn’t mention NHWC anywhere in the whole webpage.</p>
<p>I hope I have put my requirement correctly. Thanks.</p>

username: VitalyFedyunin
cooked: <p>Can you please clarify what do you mean by acceleration platform in this case. PyTorch operators (and modules) require CV tensors to be in specific indexing order NCHW. To use accelerated NHWC kernels we preserve dimensions order but laying out tensor in memory differently.</p>

username: _singh
cooked: <p>Acceleration platform is a custom processor. I am currently using the pytorch model, converting it to onnx/keras and then porting it to the processor for inference. But face this challenge of NHWC vs NHWC.</p>
<p>If possible,  can you please shed some light on the possibility of updating the model definition after permuting the layers, as mentioned in my original comment?</p>
<p>thanks!</p>

username: HermasTV
cooked: <p>any updates ? I have the same issue and I’m totally stuck</p>

username: VitalyFedyunin
cooked: <p>All PyTorch operators are written to take NCHW as dimensions order. There is no way to change it (you can only change memory format - aka how tensor laid in memory).</p>
<p>If you really want to change the order of dimensions you would need to permute each model parameter manually. Take into account that your model will not work in PyTorch anymore.</p>