如何对16GB大图像分块执行K-means聚类（无需全量加载）

核心方案：利用MiniBatchKMeans的增量训练特性 #

你之前的误区是没用到它的分批次更新聚类中心能力——不需要一次性加载全量图像，而是分块读取、逐步训练，完美适配大栅格场景。

步骤1：分块迭代训练聚类模型 #

复用你已有的栅格分块代码，每次读取固定大小的图像块，提取有效像素样本后，用partial_fit更新模型：

import rasterio
from sklearn.cluster import MiniBatchKMeans
import numpy as np

# 初始化模型：设置聚类数、批次大小（根据内存调整）
kmeans = MiniBatchKMeans(n_clusters=5, batch_size=100000, random_state=42)

# 分块读取大图像并训练
with rasterio.open("your_large_image.tif") as src:
    # 获取图像原生分块大小，也可自定义如(512,512)
    block_h, block_w = src.block_shapes[0]
    # 遍历所有图像块
    for _, window in src.block_windows():
        # 读取当前块的多波段数据
        block_data = src.read(window=window)
        # 过滤无数据像素（如果图像有nodata标识）
        mask = ~np.all(block_data == src.nodata, axis=0) if src.nodata else np.ones(block_data.shape[1:], dtype=bool)
        # 转成(样本数, 特征数)格式
        valid_samples = block_data[:, mask].T
        # 增量更新聚类中心
        kmeans.partial_fit(valid_samples)

步骤2：分块预测并输出结果 #

训练完成后，同样分块读取图像，预测每个像素的聚类标签，再写入输出栅格：

# 创建输出图像文件
with rasterio.open(
    "clustered_output.tif",
    "w",
    driver="GTiff",
    height=src.height,
    width=src.width,
    count=1,
    dtype=np.uint8,
    crs=src.crs,
    transform=src.transform,
    nodata=src.nodata
) as dst:
    for _, window in src.block_windows():
        block_data = src.read(window=window)
        mask = ~np.all(block_data == src.nodata, axis=0) if src.nodata else np.ones(block_data.shape[1:], dtype=bool)
        valid_samples = block_data[:, mask].T
        # 预测当前块的标签
        labels = kmeans.predict(valid_samples)
        # 构建输出块，填充无数据值
        output_block = np.full(block_data.shape[1:], src.nodata, dtype=np.uint8)
        output_block[mask] = labels
        # 写入当前块到输出文件
        dst.write(output_block, window=window, indexes=1)

关键优化点 #

• 批次大小调整：batch_size根据你的内存容量灵活设置，建议在10万-100万之间，平衡内存占用和训练效率。
• 无效像素过滤：提前剔除无数据、全黑等无效像素，减少不必要的计算量。
• 中心初始化优化：可先随机抽取2-3个小批次样本，用普通KMeans生成初始聚类中心，再传入MiniBatchKMeans，提升聚类稳定性。
• 避免冗余复制：直接读取原图像文件分块，完全不需要用numpy.memmap做全量复制，节省时间和磁盘空间。

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