Class: XGBoost::Classifier

Inherits:
Model
  • Object
show all
Defined in:
lib/xgboost/classifier.rb

Instance Attribute Summary

Attributes inherited from Model

#booster

Instance Method Summary collapse

Methods inherited from Model

#feature_importances, #load_model, #save_model

Constructor Details

#initialize(n_estimators: 100, objective: "binary:logistic", importance_type: "gain", **options) ⇒ Classifier

Returns a new instance of Classifier.



3
4
5
 
# File 'lib/xgboost/classifier.rb', line 3

def initialize(n_estimators: 100, objective: "binary:logistic", importance_type: "gain", **options)
  super
end

Instance Method Details

#fit(x, y, eval_set: nil, early_stopping_rounds: nil, verbose: true) ⇒ Object 



7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
 
# File 'lib/xgboost/classifier.rb', line 7

def fit(x, y, eval_set: nil, early_stopping_rounds: nil, verbose: true)
  n_classes = y.uniq.size

  params = @params.dup
  if n_classes > 2
    params[:objective] = "multi:softprob"
    params[:num_class] = n_classes
  end

  dtrain = DMatrix.new(x, label: y)
  evals = Array(eval_set).map.with_index { |v, i| [DMatrix.new(v[0], label: v[1]), "validation_#{i}"] }

  @booster = XGBoost.train(params, dtrain,
    num_boost_round: @n_estimators,
    early_stopping_rounds: early_stopping_rounds || @early_stopping_rounds,
    verbose_eval: verbose,
    evals: evals
  )
  nil
end

#predict(data) ⇒ Object 



28
29
30
31
32
33
34
35
36
37
38
39
 
# File 'lib/xgboost/classifier.rb', line 28

def predict(data)
  y_pred = super(data)

  if y_pred.first.is_a?(Array)
    # multiple classes
    y_pred.map do |v|
      v.map.with_index.max_by { |v2, _| v2 }.last
    end
  else
    y_pred.map { |v| v > 0.5 ? 1 : 0 }
  end
end

#predict_proba(data) ⇒ Object 



41
42
43
44
45
46
47
48
49
50
51
 
# File 'lib/xgboost/classifier.rb', line 41

def predict_proba(data)
  dmat = DMatrix.new(data)
  y_pred = @booster.predict(dmat)

  if y_pred.first.is_a?(Array)
    # multiple classes
    y_pred
  else
    y_pred.map { |v| [1 - v, v] }
  end
end