This example demonstrates the use of a custom F1 metric in a LightGBM implementation for binary classification with imbalanced data. The F1 score represents the harmonic mean between precision and recall, and it’s often the best measure of performance in situations with imbalanced data. The metric code is also documented here: https://github.com/dalekube/custom-metrics/blob/master/F1-score-LightGBM.R.

The ‘wisconsin’ data set contains traits about patients with cancer: https://www.rdocumentation.org/packages/imbalance/versions/1.0.0/topics/wisconsin.

library(lightgbm)
library(Laurae)
library(imbalance)

# Load the 'wisconsin' dataset
data("wisconsin")

# Massage original dataset
wi = as.data.frame(lapply(wisconsin,as.integer))
wi$Class = wi$Class-1
rm(wisconsin)

# Proportion of observations; negative vs positive outcomes
prop.table(table(wi$Class))
## 
##         0         1 
## 0.6500732 0.3499268
# Separate train/test data with 80%/20% split
n = nrow(wi)
n.idx = sample(n,n*0.80)
train.x = as.matrix(wi[n.idx,-10])
train.y = wi$Class[n.idx]
test.x = as.matrix(wi[-n.idx,-10])
test.y = wi$Class[-n.idx]

dtrain = lgb.Dataset(data=train.x,label=train.y)
dtest = lgb.Dataset(data=test.x,label=test.y)


LightGBM w/ F1 Score Metric

This implementation uses the get.max_f1() score metric from the Laurae package to compute the F1 score. This can also be used with LightGBM’s built in cross-validation function, lgb.cv().

# Define the F1 Score metric
F1_metric = function(preds,dtrain){
  
  labels = getinfo(dtrain,"label")
  F1 = get.max_f1(preds,labels)[1]
  return(list(name="F1",value=F1,higher_better=TRUE))
  
}

# Training parameters
params = list(
  objective="binary",
  boosting="gbdt",
  learning_rate=0.01,
  min_child_weight=1,
  num_leaves=5,
  max_depth=5,
  device_type="gpu",
  feature_fraction=0.5,
  bagging_fraction=0.5
)

# Train the model
lgb.model = lgb.train(
  params=params,
  data=dtrain,
  nrounds=50,
  num_threads=8,
  valids=list(test=dtest),
  early_stopping_rounds=5,
  eval=F1_metric,
  metric="None",
  verbose=1
)
## [1]: test's F1:0.886957 
## [2]: test's F1:0.9 
## [3]: test's F1:0.935484 
## [4]: test's F1:0.943089 
## [5]: test's F1:0.943089 
## [6]: test's F1:0.957265 
## [7]: test's F1:0.966102 
## [8]: test's F1:0.966667 
## [9]: test's F1:0.966102 
## [10]:    test's F1:0.966102 
## [11]:    test's F1:0.958678 
## [12]:    test's F1:0.958678 
## [13]:    test's F1:0.958678
# Print the best F1 score across the rounds
print(lgb.model$best_score)
## [1] 0.9666667