Python绘制栅格图时如何解决内存不足问题？

配备i7处理器、32GB内存及NVIDIA RTX 30显卡，在Python中绘制栅格图时遭遇内存不足问题。已做的优化包括：用rasterio的LZW算法压缩栅格、将栅格设为布尔型并裁剪至目标区域、转换为包含数组的字典。执行以下代码时触发内存错误：

for name, data in images.items():
    print(name)
    fig, ax = plt.subplots(1, 1, figsize=(7, 7))
    if '2' in name:
        plt.imshow(data, cmap='binary', vmin=0, vmax=1)
        #plt.savefig(str(name)+'IG.png',format='png',transparent=True)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        plt.show()
    else:
        im = ax.imshow(data, cmap='viridis', vmin=0, vmax=0.8)
        cbar = fig.colorbar(im, orientation='vertical', pad=0.05)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        plt.show()

错误信息：

MemoryError: Unable to allocate 7.30 GiB for an array with shape (62451, 31361) and data type float32

急需完成数据可视化，求可用库或解决方法。

1. 降采样处理栅格数组 #

直接对超大栅格做降采样，减少传入绘图函数的数据量，是最直接的内存优化方式：

import numpy as np

# 方法1：简单隔点采样（速度快）
downsampled_data = data[::10, ::10]  # 每隔10个像素取一个值，可按需调整间隔

# 方法2：均值降采样（结果更平滑）
def downsample_by_mean(arr, factor):
    return arr.reshape(arr.shape[0]//factor, factor, arr.shape[1]//factor, factor).mean(axis=(1,3))
downsampled_data = downsample_by_mean(data, 10)

降采样后再传入imshow，内存占用会按采样比例的平方级降低。

2. 改用高效可视化工具 #

GDAL直接导出图像 #

跳过Matplotlib，用GDAL直接将栅格转为图像文件，完全绕开内存渲染瓶颈：

from osgeo import gdal
import numpy as np

# 布尔型转uint8格式（0/1）
output_data = data.astype(np.uint8)

# 创建PNG输出文件
driver = gdal.GetDriverByName('PNG')
out_ds = driver.Create(f'{name}.png', output_data.shape[1], output_data.shape[0], 1, gdal.GDT_Byte)
out_band = out_ds.GetRasterBand(1)
out_band.WriteArray(output_data)

# 设置颜色映射（对应binary色板）
color_table = gdal.ColorTable()
color_table.SetColorEntry(0, (255,255,255,255))  # 白色对应0
color_table.SetColorEntry(1, (0,0,0,255))        # 黑色对应1
out_band.SetColorTable(color_table)

out_band.FlushCache()
del out_ds

Datashader处理大规模数据 #

Datashader专门为超大栅格/点云设计，会先将数据聚合到低分辨率画布再渲染，不会触发内存溢出：

import datashader as ds
from datashader.mpl_ext import dsshow
import matplotlib.pyplot as plt

fig, ax = plt.subplots(figsize=(7,7))
if '2' in name:
    dsshow(data, ds.Canvas(plot_width=700, plot_height=700), cmap='binary', vmin=0, vmax=1, ax=ax)
else:
    dsshow(data, ds.Canvas(plot_width=700, plot_height=700), cmap='viridis', vmin=0, vmax=0.8, ax=ax)

ax.set_yticklabels([])
ax.set_xticklabels([])
plt.show()

3. 优化Matplotlib内存管理 #

• 手动释放绘图资源：每次绘图后关闭当前图对象，避免内存堆积

for name, data in images.items():
    print(name)
    fig, ax = plt.subplots(1, 1, figsize=(7, 7))
    if '2' in name:
        plt.imshow(data[::10,::10], cmap='binary', vmin=0, vmax=1)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        plt.show()
        plt.close(fig)  # 释放当前图内存
    else:
        im = ax.imshow(data[::10,::10], cmap='viridis', vmin=0, vmax=0.8)
        cbar = fig.colorbar(im, orientation='vertical', pad=0.05)
        ax.set_yticklabels([])
        ax.set_xticklabels([])
        plt.show()
        plt.close(fig)

• 切换非交互式后端：如果不需要实时预览，用Agg后端减少内存开销

import matplotlib
matplotlib.use('Agg')  # 必须在导入pyplot前设置
import matplotlib.pyplot as plt

# 后续直接用plt.savefig保存图像，无需plt.show()

4. 利用GPU转移内存压力 #

你的RTX 30显卡可以承担数据处理任务，用CuPy将数组转移到显存，减少内存占用：

import cupy as cp

# 将数组移到GPU
gpu_data = cp.array(data)
# GPU上执行降采样
downsampled_gpu = gpu_data[::10, ::10]
# 转回CPU用于绘图（部分工具支持直接读取GPU数组）
downsampled_data = cp.asnumpy(downsampled_gpu)

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