Python处理大尺寸GeoTIFF栅格触发内存错误的解决方案求助

问题根因 #

你遇到的内存错误本质是操作触发了Dask分块数组的全量计算，把几十GB的栅格直接加载到内存导致的，现有代码的核心问题如下：

• 对DataArray.data直接赋值会强制触发Dask数组计算，将全量数据拉取到内存，是OOM的直接诱因
• 月度影像读取时未指定分块参数，默认全量加载进内存
• 256*256的分块尺度过小，会大幅提升Dask调度开销，也会间接增加内存占用
• reproject_match执行后已经自动对齐了坐标，手动assign_coords属于冗余操作

优化后实现代码 #

import rasterio as rio
import xarray as xr
import rioxarray
import numpy as np
import dask

# 调整调度器适配Windows环境，避免多进程开销
dask.config.set(scheduler="threads")

# 配置参数
CHUNK_SIZE = 2048  # 单块尺寸，内存不足可下调为1024
years = np.arange(2012, 2020)
monthly_file_paths = [r"\paths\to\monthly\image\tifs"]  # 路径加r避免转义
yearly_file_paths = [r"\paths\to\yearly\image\tifs"]

# 读取参考月度影像，指定分块
monthly_sample = rioxarray.open_rasterio(
    monthly_file_paths[0], 
    chunks={"x": CHUNK_SIZE, "y": CHUNK_SIZE}
).squeeze(drop=True)  # 去掉多余的波段维度，减少后续运算开销

for yearly_file in yearly_file_paths:
    # 读取年度影像，指定分块
    yearly = rioxarray.open_rasterio(
        yearly_file, 
        chunks={"x": CHUNK_SIZE, "y": CHUNK_SIZE}
    ).squeeze(drop=True)
    yearly = yearly.rio.set_nodata(0)
    
    # 裁剪+重投影匹配月度影像范围
    yearly = yearly.rio.clip_box(
        minx=monthly_sample.x.min().item(),
        miny=monthly_sample.y.min().item(),
        maxx=monthly_sample.x.max().item(),
        maxy=monthly_sample.y.max().item()
    )
    yearly_mask = yearly.rio.reproject_match(monthly_sample)
    
    # 生成二值掩膜，全程Dask延迟计算，不触发全量加载
    yearly_mask = (yearly_mask == 3).astype(np.uint8)
    # 提前持久化年掩膜到内存/磁盘，避免循环12次重复计算，内存不足可注释该行
    yearly_mask = yearly_mask.persist()

    for month_file in monthly_file_paths:
        # 读取月度影像时指定分块
        _monthly = rioxarray.open_rasterio(
            month_file, 
            chunks={"x": CHUNK_SIZE, "y": CHUNK_SIZE}
        ).squeeze(drop=True)
        # 月度影像二值化，延迟计算
        _monthly_mask = (_monthly == 2).astype(np.uint8)
        
        # 掩膜相乘，延迟计算
        yearly_monthly = yearly_mask * _monthly_mask
        yearly_monthly = yearly_monthly.rio.update_attrs(_monthly.attrs)
        
        # 写入时Dask自动逐块计算输出，不会爆内存
        yearly_monthly.rio.to_raster(
            f'{month_file}_processed.tif', 
            dtype=np.uint8,
            compress='deflate',
            tiled=True  # 输出分块TIFF，后续读取更快
        )

关键说明 #

• 全程所有运算都采用Dask延迟执行模式，只有调用to_raster时才会逐块计算写入，全流程不会加载整幅影像到内存
• 分块尺寸可根据你的设备内存调整，单块大小控制在100-500MB区间效率最优
• 如果内存非常有限，可删除yearly_mask.persist()行，牺牲一点运行速度换取更低内存占用
• 路径字符串前加r可避免Windows路径反斜杠的转义问题

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