能否为sklearn的RandomizedSearchCV设置运行时间限制？

核心思路是对基础分类器的fit方法做超时包装，在单组参数训练超过30分钟时主动终止该任务，继续执行下一组参数，不会阻塞整个参数搜索流程，同时兼容RandomizedSearchCV和GridSearchCV两种调优接口。

方案1：使用stopit第三方库实现（最简方案） #

• 先安装依赖包

  pip install stopit
  

• 定义带超时控制的分类器包装类

  from sklearn.base import BaseEstimator, ClassifierMixin
  import stopit
  
  class TimedClassifier(BaseEstimator, ClassifierMixin):
      def __init__(self, base_clf, timeout=1800):
          # timeout单位为秒，30分钟对应1800秒
          self.base_clf = base_clf
          self.timeout = timeout
          self.fit_success = True
  
      def fit(self, X, y, **fit_params):
          self.fit_success = True
          # 超时后主动终止fit进程
          with stopit.ThreadingTimeout(self.timeout) as to_ctx_mgr:
              self.base_clf.fit(X, y, **fit_params)
          if not to_ctx_mgr:
              self.fit_success = False
          return self
  
      def predict(self, X):
          if not self.fit_success:
              # 超时的模型返回默认预测值，可按需调整
              return [0]*len(X)
          return self.base_clf.predict(X)
  
      def score(self, X, y):
          if not self.fit_success:
              # 超时的参数组合返回最低分，会被自动排在结果最后
              return 0.0
          return self.base_clf.score(X, y)
  

方案2：Python内置库实现（无需额外安装依赖） #

使用multiprocessing模块实现超时控制，适配无外网安装依赖的场景：

from sklearn.base import BaseEstimator, ClassifierMixin
import multiprocessing

def _fit_worker(clf, X, y, fit_params, return_dict):
    clf.fit(X, y, **fit_params)
    return_dict["clf"] = clf

class TimedClassifier(BaseEstimator, ClassifierMixin):
    def __init__(self, base_clf, timeout=1800):
        self.base_clf = base_clf
        self.timeout = timeout
        self.fit_success = True

    def fit(self, X, y, **fit_params):
        self.fit_success = True
        manager = multiprocessing.Manager()
        return_dict = manager.dict()
        p = multiprocessing.Process(target=_fit_worker, args=(self.base_clf, X, y, fit_params, return_dict))
        p.start()
        p.join(timeout=self.timeout)
        if p.is_alive():
            p.terminate()
            p.join()
            self.fit_success = False
        else:
            self.base_clf = return_dict["clf"]
        return self

    # predict、score方法同方案1
    def predict(self, X):
        return [0]*len(X) if not self.fit_success else self.base_clf.predict(X)

    def score(self, X, y):
        return 0.0 if not self.fit_success else self.base_clf.score(X, y)

完整使用示例 #

from sklearn.model_selection import RandomizedSearchCV
from sklearn.svm import SVC
import numpy as np

# 初始化原分类器
base_clf = SVC()
# 包装为带30分钟超时的分类器
timed_clf = TimedClassifier(base_clf, timeout=1800)

# 定义参数搜索空间
param_dist = {
    "base_clf__C": np.logspace(-3, 3, 7),
    "base_clf__gamma": ["scale", "auto"]
}

# 初始化参数搜索器，原有逻辑无需修改
search = RandomizedSearchCV(timed_clf, param_dist, n_iter=10, cv=5, n_jobs=-1)
# 执行搜索，超时的参数组合会自动跳过
search.fit(X_train, y_train)

注意事项 #

• 若使用GridSearchCV，仅需将上述示例中的RandomizedSearchCV替换即可，其他逻辑完全兼容
• 超时的参数组合会在搜索结果中返回0分，可通过search.cv_results_筛选掉训练失败的参数组
• 多进程并行搜索时需将包装类放在模块顶层定义，避免序列化报错

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