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