Optuna试验结合多进程时的内存泄漏问题及解决咨询

你的内存泄漏主要来自两个核心原因：

1. 进程池未正确回收：每次Optuna试验（trial）中创建的multiprocessing.Pool没有被显式关闭销毁，导致进程持续占用内存，多次试验后内存累积耗尽。
2. Fork进程的内存继承：默认的fork进程创建方式会完整复制父进程的内存空间（包括Optuna的试验状态、历史数据、已训练模型等），每个子进程都会携带冗余内存，进一步加剧内存占用。

以下是几个可落地的修复方案，按优先级排序：

1. 使用上下文管理器确保进程池自动回收 #

用with语句包裹Pool的创建和使用，试验结束后会自动关闭进程池并回收所有子进程，避免残留内存占用。

修改后的objective函数：

def objective(trial):
    param = {
        "n_jobs": 1,  # 用整数替代字符串，避免潜在解析问题
        "num_leaves": trial.suggest_int("num_leaves", 2, 256)
    }
    lgb_params = [param for _ in range(N)]
    # 使用with上下文管理器自动管理进程池生命周期
    with multiprocessing.Pool(n_cores) as p:
        results = p.map(get_metric, zip(data, lgb_params))
    return np.mean(results)

2. 改用Spawn进程创建方式 #

避免fork带来的内存继承问题，使用spawn上下文创建进程池，子进程会从零初始化，仅携带必要的代码和数据，大幅减少内存冗余。

修改后的代码：

def objective(trial):
    param = {
        "n_jobs": 1,
        "num_leaves": trial.suggest_int("num_leaves", 2, 256)
    }
    lgb_params = [param for _ in range(N)]
    # 使用spawn上下文创建进程池
    ctx = multiprocessing.get_context('spawn')
    with ctx.Pool(n_cores) as p:
        results = p.map(get_metric, zip(data, lgb_params))
    return np.mean(results)

3. 优化数据传递减少内存拷贝 #

你的data是全局大数组，每次map时会将数据块复制到子进程内存中。可以将data设为全局变量，让子进程直接访问，避免重复拷贝；同时提前计算误差均值，减少子进程返回的数据量：

调整后的代码：

# 将data设为全局变量，避免进程间重复传递
data = [ [np.random.normal(size=(rows_per_N, cols_per_N)), np.random.normal(size=(rows_per_N, ))] for i in range(N)]

def get_metric(params):
    idx, param = params
    X, y = data[idx]  # 直接从全局数据中取对应块
    model = LGBMRegressor(**param)
    model.fit(X, y)
    return np.mean(np.abs(model.predict(X) - y))  # 提前计算均值，减少返回的数据量

def objective(trial):
    param = {
        "n_jobs": 1,
        "num_leaves": trial.suggest_int("num_leaves", 2, 256)
    }
    # 传递索引和参数，而非完整数据块
    tasks = [(i, param) for i in range(N)]
    ctx = multiprocessing.get_context('spawn')
    with ctx.Pool(n_cores) as p:
        results = p.map(get_metric, tasks)
    return np.mean(results)

4. 限制单试验内存占用 #

可以通过操作系统工具限制当前试验的内存使用，避免单试验内存溢出：

import resource

def objective(trial):
    # 限制当前进程内存使用（示例：限制为10GB）
    resource.setrlimit(resource.RLIMIT_AS, (10 * 1024**3, 10 * 1024**3))
    param = {
        "n_jobs": 1,
        "num_leaves": trial.suggest_int("num_leaves", 2, 256)
    }
    # 后续进程池代码同上

优先组合方案1+方案2+方案3，这三个改动基本能解决内存泄漏问题，同时保留多进程的速度优势。测试时可以监控工作站的内存使用情况，观察每次试验结束后内存是否能回到初始水平。

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