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

Dear PyTorch folks,

My apologies, I am a PyTorch newbie and am probably missing something basic.

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.

Here is the Colab link:

Here’s the relevant piece of code:

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 < 4)
        z_old = z_new

    return(timesteps_in_bound)

I’m not seeing any speed-up or GPU usage, though. Here is a screenshot:

Any help greatly appreciated!

Raj

You need to assign the results to a variable:

x = x.to(device='cuda')

since tensors won’t be moved inplace.

Thanks!

That was what I was missing!

Here’s my working code.

It computes fairly fast, so I think the thing limiting its speed is the matplotlib rendering in Google Colab.

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.