使用TensorFlow渲染32K曼德博集合时GPU出现OOM错误该如何解决

算力阿峰

2026-10-7

问题根因

32K分辨率（30720×17280）共计约5.3亿像素，你的代码默认使用双精度数据类型，基础张量显存占用已经接近RTX3090的24GB显存上限，额外中间计算张量申请空间时触发OOM报错，和迭代次数无关：

复数网格grid_c、迭代值current_values为complex128类型，单张张量占用 = 5.3e8 × 16字节 ≈ 8.5GB，两个合计17GB
计数张量counts为float64类型，单张占用 = 5.3e8 ×8字节 ≈4.2GB
仅三个核心张量就占用21.2GB显存，剩余显存不足以支撑计算过程中的临时变量申请，因此触发报错。

解决方案

方案1：降低数据精度（改完即可直接运行32K渲染，优先级最高）
曼德博集合常规渲染场景下float32精度完全够用，将代码中所有tf.float64替换为tf.float32，复数张量会自动变为占用8字节的complex64，核心张量总占用直接降到10.5GB左右，预留足够空间给中间计算。
修改代码示例：

# 导入tf前先加环境变量缓解显存碎片化
import os
os.environ['TF_GPU_ALLOCATOR'] = 'cuda_malloc_async'
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.float32))

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.float32)
  imag_range = tf.range(imag_start,imag_end,f,dtype=tf.float32)
  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.float32))

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

  
counts = mandelbrot(
    render_size=(30720,17280), # 32K
    center=(-0.5,0),
    zoom=4,
    cycles=200
)  
img = render(counts)
print(img.size)
img

方案2：分块渲染（适用于后续需要更高分辨率/高精度渲染场景）
将32K画布切分为多个小尺寸区块，逐个区块渲染后拼接为完整图像，单块显存占用可降至原来的1/N（N为分块数量），完全避免大张量显存占用问题。

内容的提问来源于stack exchange，提问作者TheDOC