Does torch:: nn:: Conv2d support tensors for NHWC layout in libtorch

w1005444804: <p>torch::Tensor input = torch::randn({ 1, 3, 12, 12 }, torch::kF32);<br>
input = input.permute({ 0, 2, 3, 1 });<br>
input = input.to(torch::MemoryFormat::ChannelsLast);<br>
how to set the torch::nn::Conv2d to fit the input for NHWC layout?</p>

ptrblck: <p>Call the same <code>.to(...)</code> method on the module itself.</p>

w1005444804: <aside class="quote no-group" data-username="w1005444804" data-post="1" data-topic="208193">
<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/w1005444804/48/53506_2.png" class="avatar"> w1005444804:</div>
<blockquote>
<p>torch::MemoryFormat::ChannelsLast</p>
</blockquote>
</aside>
<p><a class="mention" href="/u/ptrblck">@ptrblck</a><br>
There is no corresponding interface: .to(torch::MemoryFormat::ChannelsLast) for torch::nn::Module<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de5cae8673e7e83297e5864a648416dc2306bb12.png" data-download-href="https://discuss.pytorch.org/uploads/default/de5cae8673e7e83297e5864a648416dc2306bb12" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de5cae8673e7e83297e5864a648416dc2306bb12.png" alt="image" data-base62-sha1="vJ6zdY1EeY4da4S5BlA7egaZ05A" width="690" height="79" data-dominant-color="474561"><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">757×87 3.73 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>

ptrblck: <p>Oh, you are right. In this case, transform the <code>weight</code> directly:</p>
<pre data-code-wrap="python"><code class="lang-python">#include &lt;torch/torch.h&gt;
#include &lt;iostream&gt;

int main() {
  torch::Tensor tensor = torch::randn({1, 3, 224, 224}).to(torch::MemoryFormat::ChannelsLast).to(torch::kCUDA);
  std::cout &lt;&lt; tensor.strides() &lt;&lt; ", " &lt;&lt; tensor.sizes() &lt;&lt; std::endl;

  auto conv = torch::nn::Conv2d(torch::nn::Conv2dOptions(3, 8, 3).stride(1).bias(false));
  conv-&gt;to(torch::kCUDA);
  {
    torch::NoGradGuard no_grad;
    conv-&gt;weight = conv-&gt;weight.to(torch::MemoryFormat::ChannelsLast);
  } 

  std::cout &lt;&lt; conv-&gt;weight.strides() &lt;&lt; ", " &lt;&lt; conv-&gt;weight.sizes() &lt;&lt; std::endl;
  auto out = conv-&gt;forward(tensor);
  std::cout &lt;&lt; out.strides() &lt;&lt; ", " &lt;&lt; out.sizes() &lt;&lt; std::endl;
}
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">nsys nvprof ./main
...
 Time (%)  Total Time (ns)  Instances  Avg (ns)  Med (ns)  Min (ns)  Max (ns)  StdDev (ns)                                                  Name                                                
 --------  ---------------  ---------  --------  --------  --------  --------  -----------  ----------------------------------------------------------------------------------------------------
     85.7            23104          1   23104.0   23104.0     23104     23104          0.0  void convolve_common_engine_float_NHWC&lt;float, float, (int)128, (int)5, (int)5, (int)3, (int)3, (int…
     14.3             3841          1    3841.0    3841.0      3841      3841          0.0  void at::native::elementwise_kernel&lt;(int)128, (int)2, void at::native::gpu_kernel_impl_nocast&lt;at::n…
</code></pre>

w1005444804: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you so much <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>