为何我的CuPy KMeans代码在GPU运行时出现OutOfMemoryError？

实验基本情况 #

• 实验目标：本人为GPU编程新手，基于CuPy实现KMeans算法，在(n,3)结构的大规模数据集上运行测试，直观对比GPU与CPU运行KMeans的耗时差异
• 实验规划：设置较大聚类簇数开展对照实验
• 当前问题：调大参数规模后频繁触发内存管理错误，查阅相关资料后未形成清晰排错思路
• 运行环境：Colab Pro，配备25GB RAM资源
• 可稳定运行的参数基准：n_clusters=200、num=1000000，对应调用方式为

  run_gpu(0,200,1000000,10,True)
  

• 错误触发条件：在上述基准参数基础上进一步调大数据规模时，会弹出显存内存管理类错误提示

完整实验代码 #

import contextlib
import time

import cupy
import matplotlib.pyplot as plt
import numpy

@contextlib.contextmanager
def timer(message):
    cupy.cuda.Stream.null.synchronize()
    start = time.time()
    yield
    cupy.cuda.Stream.null.synchronize()
    end = time.time()
    print('%s:  %f sec' % (message, end - start))

    
var_kernel = cupy.ElementwiseKernel(
    'T x0, T x1, T c0, T c1', 'T out',
    'out = (x0 - c0) * (x0 - c0) + (x1 - c1) * (x1 - c1)',
    'var_kernel'
)
sum_kernel = cupy.ReductionKernel(
    'T x, S mask', 'T out',
    'mask ? x : 0',
    'a + b', 'out = a', '0',
    'sum_kernel'
)
count_kernel = cupy.ReductionKernel(
    'T mask', 'float32 out',
    'mask ? 1.0 : 0.0',
    'a + b', 'out = a', '0.0',
    'count_kernel'
)
    
    
def fit_xp(X, n_clusters, max_iter):
    assert X.ndim == 2
 
    # Get NumPy or CuPy module from the supplied array.
    xp = cupy.get_array_module(X)

    n_samples = len(X)
    
    # Make an array to store the labels indicating which cluster each sample is
    # contained.
    pred = xp.zeros(n_samples)
    
    # Choose the initial centroid for each cluster.
    initial_indexes = xp.random.choice(n_samples, n_clusters, replace=False)
    centers = X[initial_indexes]
    
    for _ in range(max_iter):
        # Compute the new label for each sample.
        distances = xp.linalg.norm(X[:, None, :] - centers[None, :, :], axis=2)
        new_pred = xp.argmin(distances, axis=1)
    
        # If the label is not changed for each sample, we suppose the
        # algorithm has converged and exit from the loop.
        if xp.all(new_pred == pred):
            break
        pred = new_pred
    
        # Compute the new centroid for each cluster.
        i = xp.arange(n_clusters)
        mask = pred == i[:, None]
        sums = xp.where(mask[:, :, None], X, 0).sum(axis=1)
        counts = xp.count_nonzero(mask, axis=1).reshape((n_clusters, 1))
        centers = sums / counts
    
    return centers, pred
    
    
def fit_custom(X, n_clusters, max_iter):
    assert X.ndim == 2
    
    n_samples = len(X)
    
    pred = cupy.zeros(n_samples,dtype='float32')
    
    initial_indexes = cupy.random.choice(n_samples, n_clusters, replace=False)
    centers = X[initial_indexes]
    
    for _ in range(max_iter):
        distances = var_kernel(X[:, None, 0], X[:, None, 1],
                                   centers[None, :, 1], centers[None, :, 0])
        new_pred = cupy.argmin(distances, axis=1)
        if cupy.all(new_pred == pred):
            break
        pred = new_pred
    
        i = cupy.arange(n_clusters)
        mask = pred == i[:, None]
        sums = sum_kernel(X, mask[:, :, None], axis=1)
        counts = count_kernel(mask, axis=1).reshape((n_clusters, 1))
        centers = sums / counts
    
    return centers, pred
    
    
def draw(X, n_clusters, centers, pred, output):
    # Plot the samples and centroids of the fitted clusters into an image file.
    for i in range(n_clusters):
        labels = X[pred == i]
        plt.scatter(labels[:, 0], labels[:, 1], c=numpy.random.rand(3))
    plt.scatter(
        centers[:, 0], centers[:, 1], s=120, marker='s', facecolors='y',
        edgecolors='k')
    plt.savefig(output)
  
    
def run_cpu(gpuid, n_clusters, num, max_iter, use_custom_kernel):
    samples = numpy.random.randn(num, 3)
    X_train = numpy.r_[samples + 1, samples - 1]
    
    with timer(' CPU '):
        centers, pred = fit_xp(X_train, n_clusters, max_iter)
    
    
    
def run_gpu(gpuid, n_clusters, num, max_iter, use_custom_kernel):
    samples = numpy.random.randn(num, 3)
    X_train = numpy.r_[samples + 1, samples - 1]
    
    with cupy.cuda.Device(gpuid):
        X_train = cupy.asarray(X_train)
    
        with timer(' GPU '):
            if use_custom_kernel:
                centers, pred = fit_custom(X_train, n_clusters, max_iter)
            else:
                centers, pred = fit_xp(X_train, n_clusters, max_iter)

排错切入点与修复方案 #

显存占用的核心来源 #

• 临时广播数组的峰值占用过高：现有代码中距离计算、掩码生成的广播操作会生成大尺寸临时数组。例如X[:, None, :] - centers[None, :, :]会生成形状为(n_samples, n_clusters, 3)的中间数组，当样本量为200万、聚类数为500时，仅这一步就会占用近12GB显存
• 分组聚合的冗余数组：mask = pred == i[:, None]会生成形状为(n_clusters, n_samples)的布尔数组，配合后续xp.where(mask[:, :, None], X, 0)的广播操作，会额外生成同维度的浮点数组，进一步推高显存峰值
• 无意义的精度冗余：默认采用float64精度存储数组，对于三维点聚类场景，float32精度完全可满足计算需求，double精度会直接让显存占用翻倍
• 缓存显存的累积占用：CuPy默认使用缓存内存分配器，循环迭代过程中产生的临时数组如果未被及时回收，会逐步累积占用显存，触发OOM错误

可落地的修复措施 #

• 分块计算距离，将样本集切分为多个小批次，每批次仅计算当前块内样本到所有聚类中心的距离，计算完成后立刻回收临时数组，避免一次性生成全量距离矩阵
• 替换高显存占用的分组聚合逻辑，取消全量广播生成mask和聚合临时数组的操作，改用逐簇遍历的聚合方式，避免生成(n_clusters, n_samples)维度的大数组
• 所有CuPy数组统一指定dtype='float32'，从源头降低一半的显存占用
• 每轮迭代结束后调用cupy.get_default_memory_pool().free_all_blocks()，主动回收不再使用的缓存显存
• 运行时实时打印各阶段的显存占用数值，精准定位显存峰值出现的代码位置

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