Custom callback to save predictions for each boosting tree

#1

My goal is to save predictions after each boosting round while I train. I’ve been following the code for the evaluation log callback and it seems I am close but I cannot figure it out.

cb.save.predictions <- function(){
  
  finalizer <- function(env){
    env$pred <- data.table(t(simplify2array(env$pred)))
    setnames(env$pred, paste0('col', 1:ncol(env$pred)))
  }
  
  callback <- function(env = parent.frame(), finalize = FALSE){
    if (is.null(env$bst)) 
      stop("'cb.save.predictions' callback requires the 'bst' booster object in its calling frame")
    
    if (finalize) 
      return(finalizer(env))
    
    if(is.null(env$pred))
      env$pred <- list()
    
    pr <- predict(env$bst, env$watchlist[[2]])
    env$pred <- c(env$pred, list(pr))
    print(length(env$pred))
  }
  
  attr(callback, "call") <- match.call()
  attr(callback, "name") <- "cb.save.predictions"
  callback
}

dt <- data.table(x = rnorm(1e3, mean = 100, sd = 10))[, y := rnorm(1e3, mean = x, sd = 10)]
xgbm_train <- xgb.DMatrix(as.matrix(dt[1:700, .(x)]), label = dt[1:700, (y)])
xgbm_test <- xgb.DMatrix(as.matrix(dt[701:1000, .(x)]), label = dt[701:1000, (y)])

model <- xgb.train(
  data = xgbm_train,
  watchlist = list(train = xgbm_train, test = xgbm_test),
  nrounds = 10,
  print_every_n = 1,
  verbose = 1,
  maximize = FALSE,
  callbacks = list(cb.save.predictions())
)

This code successfully runs but there is no pred object in the model object. What am I doing wrong? Is there an easier way to do this?

I am using xgboost_1.2.0.1

#2

I realized I was setting the pred object in the environment but was not attaching it to the final output. Updating finalizer works using:

finalizer <- function(env){
    env$pred <- data.table(t(simplify2array(env$pred)))
    setnames(env$pred, paste0('col', 1:ncol(env$pred)))
    env$bst$pred <- env$pred
  }