Spark ML中的Estimator源码解析

Spark ML中的Estimator源码解析

1. 源码加上中文注释

/**
 * :: DeveloperApi ::
 * 用于将模型拟合到数据的估计器的抽象类。
 */
@DeveloperApi
abstract class Estimator[M <: Model[M]] extends PipelineStage {

  /**
   * 使用可选参数将单个模型拟合到输入数据。
   *
   * @param dataset 输入数据集
   * @param firstParamPair 第一个参数对，覆盖嵌入参数
   * @param otherParamPairs 其他参数对。这些值会覆盖此估计器嵌入 ParamMap 中指定的任何值。
   * @return 拟合后的模型
   */
  @Since("2.0.0")
  @varargs
  def fit(dataset: Dataset[_], firstParamPair: ParamPair[_], otherParamPairs: ParamPair[_]*): M = {
    val map = new ParamMap()
      .put(firstParamPair)
      .put(otherParamPairs: _*)
    fit(dataset, map)
  }

  /**
   * 使用提供的参数映射将单个模型拟合到输入数据。
   *
   * @param dataset 输入数据集
   * @param paramMap 参数映射。这些值会覆盖此估计器嵌入 ParamMap 中指定的任何值。
   * @return 拟合后的模型
   */
  @Since("2.0.0")
  def fit(dataset: Dataset[_], paramMap: ParamMap): M = {
    copy(paramMap).fit(dataset)
  }

  /**
   * 将模型拟合到输入数据。
   */
  @Since("2.0.0")
  def fit(dataset: Dataset[_]): M

  /**
   * 使用多个参数映射将多个模型拟合到输入数据。
   * 默认实现使用 for 循环遍历每个参数映射。
   * 子类可以覆盖此方法以优化多模型训练。
   *
   * @param dataset 输入数据集
   * @param paramMaps 参数映射的数组。这些值会覆盖此估计器嵌入 ParamMap 中指定的任何值。
   * @return 拟合后的模型，与输入参数映射匹配
   */
  @Since("2.0.0")
  def fit(dataset: Dataset[_], paramMaps: Array[ParamMap]): Seq[M] = {
    paramMaps.map(fit(dataset, _))
  }

  override def copy(extra: ParamMap): Estimator[M]
}

2. 多种主要用法及其代码示例

1. 使用可选参数将单个模型拟合到输入数据：

val estimator = new MyEstimator().fit(dataset, firstParamPair, otherParamPairs)

2. 使用提供的参数映射将单个模型拟合到输入数据：

val paramMap = new ParamMap()
  .put(param1, value1)
  .put(param2, value2)
val estimator = new MyEstimator().fit(dataset, paramMap)

3. 将模型拟合到输入数据：

val estimator = new MyEstimator().fit(dataset)

4. 使用多个参数映射将多个模型拟合到输入数据：

val paramMaps = Array(paramMap1, paramMap2, paramMap3)
val estimators = new MyEstimator().fit(dataset, paramMaps)

3. 源码适用场景

Estimator是Spark ML中的抽象类，用于将模型拟合到数据。它适用于以下场景：

• 需要根据给定的数据集训练一个模型。
• 可以通过设置不同的参数映射来训练多个模型，比较它们的性能。

4. 官方链接

Apache Spark - Estimator