Avoid reshpe operation to change channel dimension

username: avish
cooked: <p>Problem: I have data in format B, a, b, c, d. (Batch size, followed by some dimensions (a, b, …)<br>
Now I want to view this data along different dimension (i.e. that dimension as first dimension) and apply some transoms in this view.<br>
For instance when, viewed in dimension of c it will be of shape B, c, -1.<br>
viewed as d it will be of shape B, d, -1.</p>
<p>My current approach is resulting in too much reshaping, which is causing very slow gradient calculation and little slow forward pass. Any hints to improve this will be helpful. Sorry for using this forum as bouncing board.</p>
<pre><code class="lang-auto">            for _ in range(K):
                shape = x.shape()
                d = random.randint(1, 4)
                #  inital = [0, 1, 2, 3, 4]
                final = [0, d, 2, 3, 4]
                final[d] = 1
                x = x.transpose(d, 1)
                x = x.reshape(size[0], size[d], -1)
                x = self.layer[d - 1](x)
                x = x.reshape([
                    size[final[0]], size[final[1]], size[final[2]],
                    size[final[3]]
                ])
                x = x.transpose(d, 1)
</code></pre>
<p>Update: It can be restated as, apply linear layer to flattened view of last k dimension of non contiguous tensor without copying it.</p>

username: anantguptadbl
cooked: <p><a class="mention" href="/u/avish">@avish</a> you should be able to use <a href="https://pytorch.org/docs/master/generated/torch.transpose.html#torch.transpose" rel="noopener nofollow ugc">torch.transpose</a></p>
<pre><code class="lang-auto">a=torch.rand(2,3,2,1,3)
b=torch.transpose(a, 3, 1)
print([a.size(), b.size()])
</code></pre>
<p>Is this not what you wanted?</p>

username: avish
cooked: <p>Yes, this is what I am currently doing. So lets say, after this we want take last 3 dimension apply flatten(2)  and pass it to nn.Linear(2*3*3, op_dim). So this flatten operation in between makes a copy of tensor (as original tensor can not be viewed because it’s stride is not compatible). This is currently causing gradient calculation too slow.</p>

username: anantguptadbl
cooked: <p><a class="mention" href="/u/avish">@avish</a> understood. That would be tricky without creating a copy of the tensor. The reason is that tranpose function leverages the same underlying storage of the original tensor</p>
<p>To illustrate the issue check the below code. After the transpose the view command gives an error but the same view command does not give an error on an “original” tensor</p>
<pre><code class="lang-auto">a=torch.rand(2,3,2,1,3)
print(a.view(2, 1, 2, -1).size())

a=torch.rand(2,1,2,3,3)
print(a.view(2, 1, 2, -1).size())

a=torch.rand(2,3,2,1,3)
b=torch.transpose(a, 3, 1)
print(b.size())
print(b.view(2, 1, 2, -1))
</code></pre>
<pre><code class="lang-auto">torch.Size([2, 1, 2, 9])
torch.Size([2, 1, 2, 9])
torch.Size([2, 1, 2, 3, 3])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-39-b6aaccdd97cd&gt; in &lt;module&gt;()
      8 b=torch.transpose(a, 3, 1)
      9 print(b.size())
---&gt; 10 print(b.view(2, 1, 2, -1))

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.
</code></pre>
<p>Will try to find a way to achieve this. There are several discuss threads and also a github thread on the same</p>

username: avish
cooked: <p>Thanks, if possible can you share those threads, I was not able to find any.<br>
For inspiration I was trying to implement something like <a href="https://github.com/google-research/vision_transformer/blob/main/vit_jax/models_mixer.py" rel="noopener nofollow ugc">mlp mixer</a> but with more dimensional data.</p>