Using randperm and arange on MPS vs CPU

username: basavyr
cooked: <p>Hey there,</p>
<p>I am performing some benchmarks on a <a href="https://arxiv.org/pdf/1409.1556" rel="noopener nofollow ugc">VGG19</a> model that was written from scratch and trained on CIFAR10 dataset.</p>
<p>The pretrained model is loaded onto GPU and then for every layer in the model some random operations should be performed on the weights. I cannot go into the exact detail of the operations and the reason for doing so, but <em>essentially</em> they consist in the following:</p>
<ul>
<li><code>torch.randperm(DIM)</code> for <strong>linear layers</strong></li>
<li><code>torch.stack([torch.stack([torch.randperm(filter_size) for _ in range(           n_filters)], dim=0) for _ in range(batch_size)], dim=0)</code> for <strong><code>conv2d</code> layers</strong></li>
</ul>
<blockquote>
<p>Here I refer to them as linear and conv2d, but I know that for VGG-like architectures, they are referred to as <em>features</em> and <em>classifiers</em>.</p>
</blockquote>
<p>Thus, in a simple <code>for</code> loop, these operations will be applied sequentially, depending on the type of the layer .</p>
<pre data-code-wrap="python"><code class="lang-python">for k, v in pretrained_vgg19.state_dict().items():
            if "weight" in k:
                if len(v.shape) == 2:
                   # linear layers
                    v.copy_(...torch.randperm()...)
                elif len(v.shape) &gt; 2:
                   # conv2d layers
                    v.copy_(...torch.stack([torch.stack([torch.randperm...)
</code></pre>
<p>Now, my <strong>issue/problem</strong> is related to the execution time of these operations: depending on the device used in<code>torch.randperm()</code> and <code>torch.arange()</code>. More precisely, if the randomization method will have the same device as the model (i.e., the GPU) computations will be slow. On the other hand, if randomization is called with <code>device="cpu"</code>, the iterations will be much faster. See the comparison in the table shown below (averaged for 10 executions in both scenarios).</p>
<div class="md-table">
<table>
<thead>
<tr>
<th>Layer</th>
<th>CPU rands (seconds)</th>
<th>GPU rands (seconds)</th>
</tr>
</thead>
<tbody>
<tr>
<td>conv1</td>
<td>0.0014</td>
<td>0.0307</td>
</tr>
<tr>
<td>conv2</td>
<td>0.0060</td>
<td>0.7935</td>
</tr>
<tr>
<td>conv3</td>
<td>0.0111</td>
<td>0.9574</td>
</tr>
<tr>
<td>conv4</td>
<td>0.0205</td>
<td>1.9211</td>
</tr>
<tr>
<td>conv5</td>
<td>0.0486</td>
<td>3.8381</td>
</tr>
<tr>
<td>conv6</td>
<td>0.0924</td>
<td>7.7111</td>
</tr>
<tr>
<td>conv7</td>
<td>0.0799</td>
<td>7.6586</td>
</tr>
<tr>
<td>conv8</td>
<td>0.0790</td>
<td>7.6429</td>
</tr>
<tr>
<td>conv9</td>
<td>0.1689</td>
<td>17.5667</td>
</tr>
<tr>
<td>conv10</td>
<td>0.3066</td>
<td>30.7506</td>
</tr>
<tr>
<td>conv11</td>
<td>0.3052</td>
<td>30.8070</td>
</tr>
<tr>
<td>conv12</td>
<td>0.3180</td>
<td>30.7441</td>
</tr>
<tr>
<td>conv13</td>
<td>0.3049</td>
<td>30.7902</td>
</tr>
<tr>
<td>conv14</td>
<td>0.3074</td>
<td>30.9785</td>
</tr>
<tr>
<td>conv15</td>
<td>0.3156</td>
<td>30.9083</td>
</tr>
<tr>
<td>conv16</td>
<td>0.3059</td>
<td>47.7401</td>
</tr>
<tr>
<td>linear17</td>
<td>0.0021</td>
<td>0.0014</td>
</tr>
<tr>
<td>linear18</td>
<td>0.1672</td>
<td>0.0001</td>
</tr>
<tr>
<td>linear19</td>
<td>0.0113</td>
<td>0.0006</td>
</tr>
</tbody>
</table>
</div><p>Shouldn’t this be the other way around? Meaning that generating new tensors via <code>torch.randperm</code> and <code>torch.arange</code> behave much faster on GPU, since the model is also loaded onto it?</p>
<p>Later edit: Forgot to mention that I am performing all benchmarks on an M3 Pro MacBook, and thus my GPU is <a href="https://pytorch.org/docs/main/notes/mps.html" rel="noopener nofollow ugc"><code>MPS</code></a>.</p>

username: Tony-Y
cooked: <p>I have performed a simple benchmark of <code>torch.randperm()</code>.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import time

for device in ['cpu', 'mps']:
  for size in [10,100,1000,10_000,100_000,1000_000]:
    since = time.time()
    for _ in range(1000):
      _ = torch.randperm(size, device=device)
    elapsed = time.time() - since
    print(f'{device} {size} {elapsed}')
</code></pre>
<p>Results on M1 Pro:</p>
<pre><code class="lang-auto">cpu 10 0.0018110275268554688
cpu 100 0.0016379356384277344
cpu 1000 0.004246711730957031
cpu 10000 0.030288219451904297
cpu 100000 0.35094308853149414
cpu 1000000 4.289053916931152
mps 10 0.1772630214691162
mps 100 0.1574409008026123
mps 1000 0.2291860580444336
mps 10000 0.34720897674560547
mps 100000 0.44776439666748047
mps 1000000 1.5657782554626465
</code></pre>
<p>So, it should be done on CPU if the size is less than 100,000.</p>
<p>You should move the final tensor on CPU to the MPS device:</p>
<pre><code class="lang-auto">import torch
import time

for device in ['cpu', 'mps']:
  since = time.time()
  _ = torch.stack([torch.randperm(10, device='cpu') for _ in range(1000)]).to(device)
  elapsed = time.time() - since
  print(f'{device} {elapsed}')
</code></pre>
<p>Results on M1 Pro:</p>
<pre><code class="lang-auto">cpu 0.0049359798431396484
mps 0.010468244552612305
</code></pre>