Reshape tensor while keeping data in a given dimension

trougnouf: <p>I would like to reshape my data, which represents a single image, such that the width and height (dims 2,3) are lowered and the batch size increases; effectively making small crops as batches out of an image.<br>
I want a specific dimension, the channels (1) to contain the same data.</p>
<p>I initially wrote a pair of functions to do something like this, but I fear it’s not backprop compatible</p>
<pre><code class="lang-auto">def img_to_batch(imgtensor, patch_size: int):
    _, ch, height, width = imgtensor.shape
    assert height%patch_size == 0 and width%patch_size == 0, 'img_to_batch: dims must be dividable by patch_size. {}%{}!=0'.format(imgtensor.shape, patch_size)
    bs = math.ceil(height/patch_size) * math.ceil(width/patch_size)
    btensor = torch.zeros([bs,ch,patch_size, patch_size], device=imgtensor.device, dtype=imgtensor.dtype)
    xstart = ystart = 0
    for i in range(bs):
        btensor[i] = imgtensor[:, :, ystart:ystart+patch_size, xstart:xstart+patch_size]
        xstart += patch_size
        if xstart+patch_size &gt; width:
            xstart = 0
            ystart += patch_size
    return btensor

def batch_to_img(btensor, height: int, width: int, ch=3):
    imgtensor = torch.zeros([1, ch, height, width], device=btensor.device, dtype=btensor.dtype)
    patch_size = btensor.shape[-1]
    xstart = ystart = 0
    for i in range(btensor.size(0)):
        imgtensor[0, :, ystart:ystart+patch_size, xstart:xstart+patch_size] = btensor[i]
        xstart += patch_size
        if xstart+patch_size &gt; width:
            xstart = 0
            ystart += patch_size
    return imgtensor
</code></pre>
<p>There is a simple view / reshape function in the torch library, but when I use it the channels do not keep their respective data. eg (I would like the data shown to contain the same elements even if out of order):</p>
<pre><code class="lang-auto">&gt;&gt;&gt; imgtens = torch.rand(1,3,4,4)
&gt;&gt;&gt; imgtens[:,0,:,:]
tensor([[[0.6830, 0.2091, 0.8786, 0.6002],
         [0.0325, 0.7217, 0.1479, 0.3478],
         [0.0880, 0.8705, 0.0929, 0.7978],
         [0.7604, 0.2658, 0.3518, 0.1969]]])
&gt;&gt;&gt; reshaped = imgtens.view([4,3,2,2])
&gt;&gt;&gt; reshaped[:,0,:,:]
tensor([[[0.6830, 0.2091],
         [0.8786, 0.6002]],

        [[0.7604, 0.2658],
         [0.3518, 0.1969]],

        [[0.3787, 0.0042],
         [0.3481, 0.2722]],

        [[0.4175, 0.8700],
         [0.1930, 0.7646]]])
</code></pre>
<p>I read that the fold/unfold function may be able to help by creating a sliding window.</p>

Nikronic: <p>Hi,</p>
<aside class="quote no-group" data-username="trougnouf" data-post="1" data-topic="85686">
<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/trougnouf/48/8006_2.png" class="avatar"> trougnouf:</div>
<blockquote>
<p>I read that the fold/unfold function may be able to help by creating a sliding window</p>
</blockquote>
</aside>
<p>Yes it works, actually, in your case you are extracting patches using sliding window then permuting channels.</p>
<p>Here is the code that will do the job for your arbitrary example:</p>
<pre><code class="lang-auto">x.unfold(2, 2, 2)[0].unfold(2, 2, 2).contiguous().view(3, -1, 2, 2).permute((1, 0, 2, 3))
</code></pre>
<p>But if you have question how it works, here is another post I have explained the calculations:</p><aside class="quote quote-modified" data-post="1" data-topic="51459">
  <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/z/aeb1de/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/change-the-dimension-of-tensor/51459">Change the dimension of tensor</a> 
  </div>
  <blockquote>
    Hi, 
I have a tensor with dimension [1, 1, 4, 6] like this: 
a = torch.tensor([[[ 1,  2,  3,  4,  5,  6],
                   [ 7,  8,  9, 10, 11, 12],
                   [13, 14, 15, 16, 17, 18],
                   [19, 20, 21, 22, 23, 24]]])

I want to change it to a tensor like this: 
[[ [[1, 2],
    [7, 8]],

   [[3, 4],
    [9, 10]],

   [[5, 6],
    [11, 12]],

   [[13, 14],
    [19, 20]],

   [[15, 16],
    [21, 22]],

   [[17, 18],
    [23, 24]] ]]
   

Is it possible? 
if we use this cod…
  </blockquote>
</aside>

<p>Bests</p>

trougnouf: <p>Thank you <a class="mention" href="/u/nikronic">@Nikronic</a> !<br>
If I understand this correctly, then the img_to_batch function I wrote above would be simply:</p>
<pre><code class="lang-auto">def img_to_batch(imgtensor, patch_size: int):
    _, ch, height, width = imgtensor.shape
    assert height%patch_size == 0 and width%patch_size == 0, (
        'img_to_batch: dims must be dividable by patch_size. {}%{}!=0'.format(
            imgtensor.shape, patch_size))
    return imgtensor.unfold(2,patch_size,patch_size).unfold(
        3,patch_size,patch_size).contiguous().view(
            ch,-1,patch_size,patch_size).permute((1,0,2,3))       
</code></pre>
<p>Note that I made a slight modification from the equivalent of “x.unfold(2, 2, 2)[0].unfold(2,” to “x.unfold(2, 2, 2).unfold(3,”, I think that it can then also handle arbitrary batch size (and just like the channels they wouldn’t get mixed up). Correct me if I’m wrong.</p>

Nikronic: <p>I do not think you can do this as <code>tensor.unfold</code> depending on some criteria creates another dimension like <code>.unsqueeze(0)</code> and because of that <code>.view</code> and <code>.permute</code> will no longer be valid due to number of dimension mismatch.</p>
<aside class="quote no-group" data-username="trougnouf" data-post="3" data-topic="85686">
<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/trougnouf/48/8006_2.png" class="avatar"> trougnouf:</div>
<blockquote>
<p>“x.unfold(2, 2, 2)[0].unfold(2,” to “x.unfold(2, 2, 2).unfold(3,”</p>
</blockquote>
</aside>
<p>Actually, I have never thought of generalizing <code>fold</code> and <code>unfold</code> method due to its tricky behavior (or at least my bad understanding).</p>

trougnouf: <p>You are right, the following test shows that data from channel 1 can be placed in channel 0</p>
<pre><code class="lang-auto">&gt;&gt;&gt; img = torch.rand(4,3,8,8)
&gt;&gt;&gt; img_to_batch(img,4)[4,0]
tensor([[0.4258, 0.3276, 0.8221, 0.6588],
        [0.5438, 0.9239, 0.0490, 0.7193],
        [0.5852, 0.7115, 0.0703, 0.1770],
        [0.4305, 0.4190, 0.2891, 0.0326]])
&gt;&gt;&gt; img[0,1]
tensor([[0.4258, 0.3276, 0.8221, 0.6588, 0.9918, 0.6219, 0.4951, 0.4356],
        [0.5438, 0.9239, 0.0490, 0.7193, 0.6819, 0.0627, 0.0361, 0.3178],
        [0.5852, 0.7115, 0.0703, 0.1770, 0.3855, 0.0666, 0.7337, 0.0240],
        [0.4305, 0.4190, 0.2891, 0.0326, 0.3457, 0.7378, 0.5640, 0.7104],
        [0.3787, 0.2371, 0.4585, 0.6150, 0.7169, 0.6518, 0.4671, 0.1212],
        [0.8061, 0.4295, 0.1194, 0.7166, 0.7526, 0.8067, 0.1612, 0.2812],
        [0.3896, 0.8208, 0.5835, 0.6830, 0.0191, 0.7138, 0.9124, 0.7285],
        [0.0963, 0.4236, 0.2779, 0.8006, 0.1528, 0.5168, 0.6543, 0.7928]])
</code></pre>

trougnouf: <p>I think I got this <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> I moved the channel axis to the front s.t. it would be safe and it passes my manual tests.</p>
<pre><code class="lang-auto">def img_to_batch(img, patch_size: int):
    _, ch, height, width = img.shape
    assert height%patch_size == 0 and width%patch_size == 0, (
        'img_to_batch: dims must be dividable by patch_size. {}%{}!=0'.format(
            img.shape, patch_size))
    assert img.dim() == 4
    return img.unfold(2, patch_size, patch_size).unfold(
        3, patch_size, patch_size).transpose(1,0).reshape(
            ch, -1, patch_size, patch_size).transpose(1,0)
</code></pre>