OOM on GPU with tensorflow while making Mandelbrot

Question

Trying to get a beautiful Mandelbrot picture, the code works great with 16k resolution but I can't get it to render a 32k image I tried lowering cycles to 50 but still no difference Specs - i9 10900k & RTX 3090 24gb I get an OOM message saying

W tensorflow/core/common_runtime/bfc_allocator.cc:456] Allocator (GPU_0_bfc) ran out of memory trying to allocate 7.91GiB (rounded to 8493465600)requested by op Mul If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation.

here is my code

import tensorflow as tf
import numpy as np

import PIL.Image
from io import BytesIO
from IPython.display import Image, display

def render(a):
  a_cyclic = (a*0.3).reshape(list(a.shape)+[1])
  img = np.concatenate([10+20*np.cos(a_cyclic),
                        30+50*np.sin(a_cyclic),
                        155-80*np.cos(a_cyclic)], 2)
  img[a==a.max()] = 0
  a = img
  a = np.uint8(np.clip(a, 0, 255))
  f = BytesIO()
  return PIL.Image.fromarray(a)

#@tf.function
def mandelbrot_helper(grid_c, current_values, counts,cycles):
  
  for i in range(cycles):
    temp = current_values*current_values + grid_c
    not_diverged = tf.abs(temp) < 4
    current_values.assign(temp),
    counts.assign_add(tf.cast(not_diverged, tf.float64))

def mandelbrot(render_size,center,zoom,cycles):
  f = zoom/render_size[0]
  real_start = center[0]-(render_size[0]/2)*f
  real_end = real_start + render_size[0]*f 
  imag_start = center[1]-(render_size[1]/2)*f
  imag_end = imag_start + render_size[1]*f 

  real_range = tf.range(real_start,real_end,f,dtype=tf.float64)
  imag_range = tf.range(imag_start,imag_end,f,dtype=tf.float64)
  real, imag = tf.meshgrid(real_range,imag_range)
  grid_c = tf.constant(tf.complex(real, imag))
  current_values = tf.Variable(grid_c)
  counts = tf.Variable(tf.zeros_like(grid_c, tf.float64))

  mandelbrot_helper(grid_c, current_values,counts,cycles)
  return counts.numpy()

  
counts = mandelbrot(
    render_size=(30720,17280), # 32K
    #render_size=(15360,8640), # 16K
    #render_size=(7680,4320), # 8K
    #render_size=(3840,2160), # 4K
    #render_size=(1920,1080), # HD
    center=(-0.5,0),
    zoom=4,
    cycles=200
)  
img = render(counts)
print(img.size)
img

#img.save("E:/Python/Python3/TestingSO/Images/test.png")

OOM on GPU with tensorflow while making Mandelbrot

0 Answers0