Q: Trying to run Mandelbrot on GPU via PyTorch, but not seeing speed-up

username: Rajeev_Raizada
cooked: <p>Dear PyTorch folks,</p>
<p>My apologies, I am a PyTorch newbie and am probably missing something basic.</p>
<p>I made a click-to-zoom Mandelbrot set in Google Colab, and am trying to transfer it to Torch, in order to get GPU speed-up. I have selected a GPU runtime and it is showing up in nvidia-smi, and I am passing the tensor to CUDA.</p>
<p>Here is the Colab link:</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/17xHEFdUVG8TCgme8DSQP--ZLo7Wsoe6L?usp=sharing">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/16072945715e2f0a5cb0d85e8ab66893/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/17xHEFdUVG8TCgme8DSQP--ZLo7Wsoe6L?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/17xHEFdUVG8TCgme8DSQP--ZLo7Wsoe6L?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colab</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Here’s the relevant piece of code:</p>
<pre><code class="lang-auto">def raj_mandel_torch(xmin,xmax,ymin,ymax,num_grid_points,max_its):
    real_range = np.linspace(xmin,xmax,num_grid_points)
    imag_range = np.linspace(ymin,ymax,num_grid_points)
    timesteps_in_bound = np.zeros((num_grid_points,num_grid_points))
    real_grid, imag_grid = np.meshgrid(real_range, imag_range)
    complex_array = real_grid + imag_grid * 1j
    
    complex_tensor = torch.from_numpy(complex_array)
    z_old = torch.zeros(num_grid_points, num_grid_points)
    z_new = torch.zeros(num_grid_points, num_grid_points)
    abs_val_sq = torch.zeros(num_grid_points, num_grid_points)
    timesteps_in_bound = torch.zeros(num_grid_points, num_grid_points)
 
    complex_tensor.to(device='cuda')
    z_old.to(device='cuda')
    z_new.to(device='cuda')
    abs_val_sq.to(device='cuda')
    timesteps_in_bound.to(device='cuda')

    for this_iter_num in range(0,max_its):
        z_new = z_old**2 + complex_tensor
        abs_val_sq = torch.real(z_new)**2 + torch.imag(z_new)**2
        timesteps_in_bound += (abs_val_sq &lt; 4)
        z_old = z_new

    return(timesteps_in_bound)
</code></pre>
<p>I’m not seeing any speed-up or GPU usage, though. Here is a screenshot:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bb7f1190813ce66ca53039c469cf71a315b39d5e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/bb7f1190813ce66ca53039c469cf71a315b39d5e" title="Screenshot 2024-11-27 at 8.21.17 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb7f1190813ce66ca53039c469cf71a315b39d5e_2_690x453.jpeg" alt="Screenshot 2024-11-27 at 8.21.17 AM" data-base62-sha1="qKFwS9Mo1nvhUXf8aYyazC0rPqu" width="690" height="453" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb7f1190813ce66ca53039c469cf71a315b39d5e_2_690x453.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb7f1190813ce66ca53039c469cf71a315b39d5e_2_1035x679.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bb7f1190813ce66ca53039c469cf71a315b39d5e_2_1380x906.jpeg 2x" data-dominant-color="E8DEEB"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2024-11-27 at 8.21.17 AM</span><span class="informations">1920×1263 147 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Any help greatly appreciated!</p>
<p>Raj</p>

username: lennarteingunia
cooked: <aside class="quote" data-post="1" data-topic="39318">
  <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/intel_novel/48/9466_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/moving-tensor-to-cuda/39318">Moving tensor to cuda</a> 
  </div>
  <blockquote>
    Hi, this works, a = torch.LongTensor(1).random_(0, 10).to("cuda"). but this won’t work: 
a = torch.LongTensor(1).random_(0, 10)
a.to(device="cuda")

Is this per design, maybe I am simple missing something to convert tensor from CPU to CUDA?
  </blockquote>
</aside>


username: ptrblck
cooked: <aside class="quote no-group" data-username="Rajeev_Raizada" data-post="1" data-topic="213528">
<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/rajeev_raizada/48/74226_2.png" class="avatar"> Rajeev_Raizada:</div>
<blockquote>
<pre><code class="lang-auto">    complex_tensor.to(device='cuda')
    z_old.to(device='cuda')
    z_new.to(device='cuda')
    abs_val_sq.to(device='cuda')
    timesteps_in_bound.to(device='cuda')
</code></pre>
</blockquote>
</aside>
<p>You need to assign the results to a variable:</p>
<pre data-code-wrap="python"><code class="lang-python">x = x.to(device='cuda')
</code></pre>
<p>since tensors won’t be moved inplace.</p>

username: Rajeev_Raizada
cooked: <p>Thanks!</p>
<p>That was what I was missing!</p>

username: Rajeev_Raizada
cooked: <p>Here’s my working code.</p><aside class="onebox allowlistedgeneric" data-onebox-src="https://colab.research.google.com/drive/1805c-3GPVreVFcvlUFnqdgiAAYV7KkKp?usp=sharing">
  <header class="source">
      <img src="https://ssl.gstatic.com/colaboratory-static/common/16072945715e2f0a5cb0d85e8ab66893/img/favicon.ico" class="site-icon" width="128" height="128">

      <a href="https://colab.research.google.com/drive/1805c-3GPVreVFcvlUFnqdgiAAYV7KkKp?usp=sharing" target="_blank" rel="noopener nofollow ugc">colab.research.google.com</a>
  </header>

  <article class="onebox-body">
    <img width="260" height="260" src="https://colab.research.google.com/img/colab_favicon_256px.png" class="thumbnail onebox-avatar">

<h3><a href="https://colab.research.google.com/drive/1805c-3GPVreVFcvlUFnqdgiAAYV7KkKp?usp=sharing" target="_blank" rel="noopener nofollow ugc">Google Colab</a></h3>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>It computes fairly fast, so I think the thing limiting its speed is the matplotlib rendering in Google Colab.</p>
<p>I’m sure that more modern plotting libraries would be faster (plotly, bokey, seaborn etc.), but I haven’t been able to find one that supports getting the coords from a mouse-click in Google Colab. That’s mostly due to restrictions from Colab, but I don’t know of another free and shareable platform that allows GPU access. I tred Kaggle but it didn’t render the plots at all.</p>