How to fold the unfolded?

username: RoySadaka
cooked: <p>EDITED: added channels (making things even harder <img src="https://discuss.pytorch.org/images/emoji/apple/worried.png?v=12" title=":worried:" class="emoji" alt=":worried:" loading="lazy" width="20" height="20">)</p>
<p>Hey there,<br>
I have a tensor of shape [batch, C, H, W] which is a batch of images with 3 channels</p>
<pre><code class="lang-auto">batch = torch.rand((2,3,64,64))
</code></pre>
<p>I’m manipulating this batch, and my goal is to return to this shape after the manipulation was done.<br>
The manipulation is to slice each image to 16x16 squares and represent each square as a flattened vector, and lastly, since each image was sliced, we’ll take all the slices and represent them as a sequence.</p>
<p>The manipulation is as follow:</p>
<pre><code class="lang-auto">patched_squares = batch.unfold(2,16,16).unfold(3,16,16)
</code></pre>
<p>which should slice each image-channel to 16x16 squares, so the output shape will be:</p>
<blockquote>
<p>torch.Size([2, 3, 4, 4, 16, 16])</p>
</blockquote>
<p>meaning, a batch of size 2, 3 channels, with 4 rows and 4 columns that consist of 16x16 squares</p>
<p>next I’m flattening the 16x16 squares to a vector, like so:</p>
<pre><code class="lang-auto">patched_squares = patched_squares.flatten(start_dim=4, end_dim=5)
</code></pre>
<p>now the size is:</p>
<blockquote>
<p>torch.Size([2, 3, 4, 4, 256])</p>
</blockquote>
<p>meaning, a batch of size 2, 3 channels, with 4 rows and 4 columns that consist of 256 flattened vector</p>
<p>Next, I’m flattening the 4x4 to be a sequence:</p>
<pre><code class="lang-auto">patched_squares = patched_squares.flatten(start_dim=2, end_dim=3)
</code></pre>
<p>which gives a batch of size 2, 3 channels, with sequence length = 16, where each sample in the sequence is a 256 vector</p>
<blockquote>
<p>torch.Size([2, 3, 16, 256])</p>
</blockquote>
<p>Lastly, I’m concatenating the channels of each “square” like so:</p>
<pre><code class="lang-auto">patched_squares = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in patched_squares])
</code></pre>
<blockquote>
<p>torch.Size([2, 16, 768])</p>
</blockquote>
<p>The full manipulation code is:</p>
<pre><code class="lang-auto">batch = torch.rand((2,3,64,64))
patched_squares = batch.unfold(2,16,16).unfold(3,16,16).flatten(start_dim=4, end_dim=5).flatten(start_dim=2, end_dim=3)
patched_squares = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in patched_squares])
</code></pre>
<p>And now to the challenge itself <img src="https://discuss.pytorch.org/images/emoji/apple/upside_down_face.png?v=12" title=":upside_down_face:" class="emoji" alt=":upside_down_face:" loading="lazy" width="20" height="20"><br>
After running some computation on <code>patched_squares</code>, I would like to “refold” it back to how it was,<br>
so from [2, 16, 768] back to [2, 3, 64, 64], while maintaining the actual original positions.</p>
<p>I’m having a hard time to refold, I would appreciate any help in this manner.</p>
<p>Roy</p>

username: RoySadaka
cooked: <p>For future reference, iv’e managed to re-fold it, here is the solution I came up with:</p>
<pre><code class="lang-auto">#forward
start = torch.rand((2,3,64,64))
forward1 = start.unfold(2,16,16)
forward2 = forward1.unfold(3,16,16)
forward3 = forward2.flatten(start_dim=4, end_dim=5)
forward4 = forward3.flatten(start_dim=2, end_dim=3)
forward5 = torch.stack([torch.cat([p[0],p[1],p[2]], dim=1) for p in forward4])


# going back
back1 = forward5.unfold(2,256,256).permute(0,2,1,3) 	# back1 == forward4
back2 = back1.unfold(2,4,4).permute(0,1,2,4,3) 			# back2 == forward3
back3 = back2.unfold(4,16,16) 							# back3 == forward2
back4 = back3.permute(0,1,2,3,5,4).reshape(2,3,4,-1,16) # back4 == forward1
back5 = back4.permute(0,1,2,4,3).reshape(2,3,-1,64)		# back5 == start

(back5 == start).all()
</code></pre>
<p>I don’t know if it is the most optimal way to do it, but it works <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>