Apply vectorized random permutations to conv2d weight tensors

username: basavyr
cooked: <p>I want to have a method that will take the weight of a <code>conv2d</code> layer, typically in the shape <code>C_out, C_in, H, W</code>, where <code>C_out</code> represents the output channels, <code>C_in</code> represents the input channels, while <code>H</code> and <code>W</code> are the dimensions of filters (kernels).</p>
<p>As an example, consider a weight:</p>
<pre data-code-wrap="python"><code class="lang-python">w = torch.randint(0, 5, (4, 1, 3, 3))
tensor([[[[3, 3, 1],
          [3, 4, 2],
          [1, 2, 3]]],


        [[[3, 0, 1],
          [3, 1, 3],
          [0, 2, 2]]],


        [[[4, 0, 1],
          [2, 4, 3],
          [2, 4, 0]]],


        [[[0, 0, 4],
          [2, 2, 1],
          [3, 4, 0]]]])
</code></pre>
<p>For this weight, I would like to apply <strong>line-wise permutations</strong> for every filter. This means taking each of the four filters, which can be considered as a <code>3x3</code> matrix and randomly permute the lines.</p>
<p>I can create this implementation but only through a <em>naive</em> <code>for</code> loop:</p>
<pre data-code-wrap="python"><code class="lang-python">def test_conv():
    w = torch.randint(0, 5, (4, 1, 3, 3))

    w_permuted = w.clone()
    for idx, kernel in enumerate(w_permuted):
        # generate random permutation
        key = torch.randperm(kernel.shape[1])
        # permute lines of the current 3x3 kernel with generated indices
        w_permuted[idx] = kernel[:, key]

    print(w)
    print(w_permuted)
</code></pre>
<p><strong>Q:</strong> Is there a more elegant approach using PyTorch to generate <code>w_permuted</code> in a vectorized approach?</p>
<p>Thanks in advance</p>

username: ptrblck
cooked: <p>If I understand your use case correctly you want to permute the rows of all kernels randomly without mixing them between kernels.<br>
I think your code is fine as it’s clearly showing your implementation and I don’t know if this is currently a bottleneck in your code.<br>
However, you could try to use <code>torch.gather</code> as seen here and compare it against your approach:</p>
<pre data-code-wrap="python"><code class="lang-python">w = torch.arange(4*2*3*3).view(4, 2, 3, 3).to(device)
w_ = w.clone()
key = torch.stack([torch.randperm(w_.shape[2]) for _ in range(w_.shape[0] * w_.shape[1])]).to(device)
torch.gather(w_, 2, key.view(4, 2, 3, 1).expand_as(w_))
</code></pre>

username: KFrank
cooked: <p>Hi Robert!</p>
<aside class="quote no-group quote-modified" data-username="basavyr" data-post="1" data-topic="219841" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/basavyr/48/73507_2.png" class="avatar"> basavyr:</div>
<blockquote>
<p>For this weight, I would like to apply <strong>line-wise permutations</strong> for every filter. This means taking each of the four filters, which can be considered as a <code>3x3</code> matrix and randomly permute the lines.</p>
<p>I can create this implementation but only through a <em>naive</em> <code>for</code> loop:<br>
…<br>
<strong>Q:</strong> Is there a more elegant approach using PyTorch to generate <code>w_permuted</code> in a vectorized approach?</p>
</blockquote>
</aside>
<p>Yes.  The largest inefficiency in your approach would be the python loop.  This can be<br>
avoided by using <code>argsort()</code> to compute a batch of permutations without any python loop.</p>
<p>Please see this post:</p>
<aside class="quote quote-modified" data-post="2" data-topic="184339">
  <div class="title">
    <div class="quote-controls"></div>
    <img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/48.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/batched-permutation/184339/2">Batched permutation</a> 
  </div>
  <blockquote>
    Hi Javier! 

The problem is that pytorch (still) does not offer a batch version of 
randperm().  The standard work-around is to compute argsort() 
for a batch of random vectors. 
Thus: 
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; tensor_to_permute = torch.randint (100, (3, 4, 4))   # some test data
&gt;&gt;&gt; perms = torch.rand (3, 4).argsort (dim = 1)          # a way to get a batch of permutations
&gt;&gt;&gt;
&gt;&gt;&gt; permuted_tensor = torch.empty_like (tensor_to_permute)
&gt;&gt;&gt; for i in range(tensor_to_permute.shape[0…
  </blockquote>
</aside>

<p>As noted in the linked post, <code>argsort()</code> has <code>O (n log (n))</code> cost, with a factor of <code>log (n)</code><br>
that is theoretically unnecessary when generating permutations.  But for your example, your<br>
<code>n</code> is <code>3</code>, so the asymptotic inefficiency is fully irrelevant (as it will be for an <code>n</code> coming from<br>
any reasonable convolution kernel).</p>
<p>Best.</p>
<p>K. Frank</p>

username: basavyr
cooked: <p>Thanks a lot for providing this</p>

username: basavyr
cooked: <p>Interesting proposal!!! I was not aware of <code>argsort()</code></p>