在H2O中使用XGBoost时OneHotEncoding引发Java堆内存错误

在H2O中运行XGBoost训练时触发Java堆空间错误，手动执行del df、del something、import gc、gc.collect()清理内存后可正常运行。相关代码及错误栈如下：

训练代码 #

import h2o
from h2o.tree import H2OTree
from h2o.estimators import H2OIsolationForestEstimator, H2OXGBoostEstimator

encoding = "one_hot_explicit"

baseModel = H2OXGBoostEstimator(model_id = modelId, ntrees = 100, 
                                max_depth = 3,seed = 0xDECAF,
                                sample_rate = 1,
                                categorical_encoding = encoding,
                                keep_cross_validation_predictions=True, 
                                nfolds = 10
                                )
# TRAIN DATA
baseModel.train(x = predictor_columns, y = "label", training_frame = train.rbind(valid))

错误栈信息 #

Traceback (most recent call last):
      File "/docs/code/000_pyGraph/dec_rf_gb_xgb.py", line 151, in <module>
        decxgb.xgb_cvs(df=df, year=year, model_path=model_path, 
      File "/docs/code/000_pyGraph/dec_xgb.py", line 90, in xgb_cvs
        baseModel.train(x = predictor_columns, y = "label", training_frame = train.rbind(valid))
      File "/home/miniconda3/envs/tf-gpu-mem-day/lib/python3.10/site-packages/h2o/estimators/estimator_base.py", line 123, in train
        self._train(parms, verbose=verbose)
      File "/home/miniconda3/envs/tf-gpu-mem-day/lib/python3.10/site-packages/h2o/estimators/estimator_base.py", line 215, in _train
        job.poll(poll_updates=self._print_model_scoring_history if verbose else None)
      File "/home/miniconda3/envs/tf-gpu-mem-day/lib/python3.10/site-packages/h2o/job.py", line 90, in poll
        raise EnvironmentError("Job with key {} failed with an exception: {}\nstacktrace: "
    OSError: Job with key $03017f00000132d4ffffffff$_8508e7043b6647f7868aa83a3f6842d4 failed with an exception: DistributedException from /127.0.0.1:54321: 'Java heap space', caused by java.lang.OutOfMemoryError: Java heap space
    stacktrace: 
    DistributedException from /127.0.0.1:54321: 'Java heap space', caused by java.lang.OutOfMemoryError: Java heap space
        at water.MRTask.getResult(MRTask.java:660)
        at water.MRTask.getResult(MRTask.java:670)
        at water.MRTask.doAll(MRTask.java:530)
        at water.MRTask.doAll(MRTask.java:412)
        at water.MRTask.doAll(MRTask.java:397)
        at water.fvec.Vec.doCopy(Vec.java:514)
        at water.fvec.Vec.makeCopy(Vec.java:500)
        at water.fvec.Vec.makeCopy(Vec.java:493)
        at water.fvec.Vec.makeCopy(Vec.java:487)
        at water.util.FrameUtils$CategoricalOneHotEncoder$CategoricalOneHotEncoderDriver.compute2(FrameUtils.java:768)
        at water.H2O$H2OCountedCompleter.compute(H2O.java:1677)
        at jsr166y.CountedCompleter.exec(CountedCompleter.java:468)
        at jsr166y.ForkJoinTask.doExec(ForkJoinTask.java:263)
        at jsr166y.ForkJoinPool$WorkQueue.runTask(ForkJoinPool.java:976)
        at jsr166y.ForkJoinPool.runWorker(ForkJoinPool.java:1479)
        at jsr166y.ForkJoinWorkerThread.run(ForkJoinWorkerThread.java:104)
    Caused by: java.lang.OutOfMemoryError: Java heap space

解决方案 #

• 调整H2O的Java堆内存上限
  启动H2O时直接分配更大的堆内存，根据机器可用内存调整（比如16G、32G）：

  h2o.init(max_mem_size="16G")
  

  若通过命令行启动H2O，使用：

  java -Xmx16g -jar h2o.jar
  

• 优化类别特征编码方式
  one_hot_explicit会为高基数类别特征生成大量稀疏列，大幅占用内存。替换为更高效的编码方式：

  encoding = "label_encoding"  # 适合基数较高的特征
  # 或针对分类任务使用目标编码
  encoding = "target_encoding"
  

  同时提前过滤掉低重要性、高基数的类别特征，减少总特征数量。

• 降低交叉验证的内存开销
  当前nfolds=10且keep_cross_validation_predictions=True会保存大量交叉验证预测结果，可做以下调整：

  • 减少折数，比如设置nfolds=5
  • 若不需要后续使用交叉验证预测结果，关闭该参数：keep_cross_validation_predictions=False

• 自动整合内存清理逻辑
  把手动清理步骤整合到训练前，避免重复手动操作：

  # 删除Python侧无用变量
  del df, something  # 替换为实际不再使用的变量名
  import gc
  gc.collect()
  # 清理H2O中未被引用的Frame
  h2o.remove_all()
  

• 数据集过大时的分批处理
  若数据集远超内存容量，可使用H2O的流式加载方式逐步导入数据，优先结合前几种方案优化，XGBoost在H2O中以全量训练为主，此方案为备选。

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