Python使用joblib与multiprocessing实现数据分块并行处理

解决单线程性能瓶颈

当面对大量文件或超大内存数组（如遥感影像波段计算）时，传统的for顺序遍历会引发严重的I/O与CPU等待问题。Python受全局解释器锁（GIL）限制，纯多线程难以真正突破多核算力上限。因此，实际工程中常采用基于独立进程的multiprocessing模块或专为高负载计算的joblib库进行分块并行计算，以显著缩短流水线总耗时。

环境准备与依赖说明

部分网络资料误将标准库列为需手动安装的第三方包，请以官方文档为准：

• multiprocessing：Python内置标准库，开箱即用，跨平台支持多进程通信与同步。
• joblib：推荐安装pip install joblib，底层优化了磁盘与内存缓存机制，API设计更贴合科学计算场景。
• gdal 与 numpy：地理空间数据解析依赖，请确保通过conda或pip安装了匹配当前Python版本的发行版。

核心工具类与计算逻辑定义

在并行任务下发前，需封装稳定的文件读写接口，避免每个子进程重复初始化GDAL驱动或消耗额外启动资源。

import numpy as np
from osgeo import gdal
import time
from multiprocessing import cpu_count, Pool
from joblib import Parallel, delayed

class GdalUtil:
    @staticmethod
    def read_file(raster_file, read_band=None):
        """读取栅格数据"""
        gdal.AllRegister()
        gdal.SetConfigOption('gdal_FILENAME_IS_UTF8', 'YES')
        dataset = gdal.Open(raster_file, gdal.GA_ReadOnly)
        if dataset is None:
            print(f'打开图像 {raster_file} 失败.')
            return None
        
        rw, rh = dataset.RasterXSize, dataset.RasterYSize
        if read_band is None:
            return dataset.ReadAsArray(0, 0, rw, rh)
        return dataset.GetRasterBand(read_band).ReadAsArray(0, 0, rw, rh)

    @staticmethod
    def read_block_data(dataset, band_num, cols_read, rows_read, start_col=0, start_row=0):
        return dataset.GetRasterBand(band_num).ReadAsArray(start_col, start_row, cols_read, rows_read)

    @staticmethod
    def get_raster_band(raster_path):
        gdal.AllRegister()
        gdal.SetConfigOption('gdal_FILENAME_IS_UTF8', 'YES')
        ds = gdal.Open(raster_path, gdal.GA_ReadOnly)
        return ds.RasterCount if ds else 0

    @staticmethod
    def get_file_size(raster_path):
        gdal.AllRegister()
        gdal.SetConfigOption('gdal_FILENAME_IS_UTF8', 'YES')
        ds = gdal.Open(raster_path, gdal.GA_ReadOnly)
        return (ds.RasterXSize, ds.RasterYSize) if ds else (0, 0)

    @staticmethod
    def get_file_geotransform(raster_path):
        gdal.AllRegister()
        gdal.SetConfigOption('gdal_FILENAME_IS_UTF8', 'YES')
        ds = gdal.Open(raster_path, gdal.GA_ReadOnly)
        return ds.GetGeoTransform() if ds else []

    @staticmethod
    def get_file_proj(raster_path):
        gdal.AllRegister()
        gdal.SetConfigOption('gdal_FILENAME_IS_UTF8', 'YES')
        ds = gdal.Open(raster_path, gdal.GA_ReadOnly)
        return ds.GetProjection() if ds else ''

    @staticmethod
    def write_file(dataset, geotransform, project, output_path, out_format='GTiff', eType=gdal.GDT_Float32):
        """写入栅格"""
        if np.ndim(dataset) == 3:
            out_band, out_rows, out_cols = dataset.shape
        else:
            out_band, out_rows, out_cols = 1, dataset.shape[0], dataset.shape[1]

        out_driver = gdal.GetDriverByName(out_format)
        if out_driver is None:
            print(f'格式 {out_format} 不支持 Create() 方法.')
            return

        out_ds = out_driver.Create(output_path, xsize=out_cols, ysize=out_rows, bands=out_band, eType=eType)
        out_ds.SetGeoTransform(geotransform)
        out_ds.SetProjection(project)

        if out_band == 1:
            out_ds.GetRasterBand(1).WriteArray(dataset)
        else:
            for i in range(out_band):
                out_ds.GetRasterBand(i + 1).WriteArray(dataset[i])
        del out_ds

def cal_ndvi(task_params):
    """
    计算NDVI
    :param task_params: [文件路径, 起始列, 起始行, 宽度步长, 高度步长]
    :return: NDVI局部数组
    """
    input_file, start_col, start_row, cols_step, rows_step = task_params
    dataset = gdal.Open(input_file, gdal.GA_ReadOnly)
    if dataset is None: return np.zeros((rows_step, cols_step))
    
    nir = GdalUtil.read_block_data(dataset, 4, cols_step, rows_step, start_col, start_row)
    red = GdalUtil.read_block_data(dataset, 3, cols_step, rows_step, start_col, start_row)
    
    ndvi = (nir - red) / (nir + red)
    ndvi[(ndvi > 1.5) | (ndvi < -1)] = 0
    return ndvi

分块并行处理主流程

第1步：配置路径并划分网格

目的：将巨型矩阵切割为独立小块，使多个CPU核心可互不干扰地同时处理不同空间区域。

# 基础配置
input_file = r'D:\originalData\GF1\namucuo2021.tif'
output_file = r'D:\originalData\GF1\namucuo2021_ndvi.tif'
method = 'joblib'  # 可选: 'joblib' 或 'multiprocessing'

raster_cols, raster_rows = GdalUtil.get_file_size(input_file)
geotransform = GdalUtil.get_file_geotransform(input_file)
project = GdalUtil.get_file_proj(input_file)

rows_block_size = 50
cols_block_size = 50
tasks = []

for j in range(0, raster_rows, rows_block_size):
    for i in range(0, raster_cols, cols_block_size):
        rows_step = rows_block_size if j + rows_block_size < raster_rows else raster_rows - j
        cols_step = cols_block_size if i + cols_block_size < raster_cols else raster_cols - i
        tasks.append([input_file, i, j, cols_step, rows_step])

第2步：启动进程池或Joblib执行计算

目的：利用多核能力并发执行cal_ndvi函数。joblib语法极简且内置底层优化；multiprocessing.Pool为原生实现，适合重度系统级控制。

t1 = time.time()
if method == 'multiprocessing':
    pool = Pool(processes=cpu_count() - 1)
    results = pool.map(cal_ndvi, tasks)
    pool.close()
    pool.join()
else:
    results = Parallel(n_jobs=-1)(delayed(cal_ndvi)(item) for item in tasks)
t2 = time.time()
print("Total time:", t2 - t1)

第3步：结果拼接与落盘导出

目的：将各核心返回的碎片数组按原始坐标映射回零矩阵，并完成栅格元数据绑定与写入。

out_data = np.zeros([raster_rows, raster_cols], dtype='float')
for result, params in zip(results, tasks):
    sc, sr, cs, rs = params[1], params[2], params[3], params[4]
    out_data[sr:sr + rs, sc:sc + cs] = result

GdalUtil.write_file(out_data, geotransform, project, output_file)

注意事项与排错指南

嵌套循环并发冲突：在使用双重for遍历多级目录时，若在内外层同时调用多进程会导致子进程递归创建从而引发系统死锁。此时建议外层使用joblib调度整体批次，内层回归普通循环或单次multiprocessing提交，可彻底规避此问题。

序列化与后端选择：joblib在首次调用时会尝试生成缓存文件，若遇权限拒绝或磁盘满载，可显式传递memory=None跳过缓存。Windows环境下若遇到Pickle报错，建议在Parallel中指定backend='loky'以提升兼容性。

简要总结

通过合理的网格切分策略与进程级并发框架，Python能够突破单核算力天花板，成倍提升栅格影像批处理或大规模张量运算的效率。掌握索引回卷对齐与句柄生命周期管理，是保障生产环境稳定运行的核心。以上就是Python使用joblib与multiprocessing实现数据分块并行处理的详细内容，更多关于Python高性能计算与地理空间数据工程实践的资料请关注本站其它相关文章！