external_memory.c

 
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <xgboost/c_api.h>
 
#define safe_xgboost(err) \
if ((err) != 0) { \
 fprintf(stderr, "%s:%d: error in %s: %s\n", __FILE__, __LINE__, #err, \
XGBGetLastError()); \
exit(1); \
  }
 
#define N_BATCHS 32
#define BATCH_LEN 512
 
/* 缩写 */
typedef DMatrixHandle DMatrix;
typedef BoosterHandle Booster;
 
typedef struct _DataIter {
  /* 每批数据 */
  float **data;
  /* 每批标签 */
  float **labels;
  /* 每批长度 */
  size_t *lengths;
  /* 总批数 */
  size_t n;
  /* 当前迭代 */
  size_t cur_it;
 
  /* 私有字段 */
DMatrix _proxy;
  char _array[128];
} DataIter;
 
#define safe_malloc(ptr) \
if ((ptr) == NULL) { \
 fprintf(stderr, "%s:%d: 分配内存失败.\n", __FILE__, \
__LINE__); \
exit(1); \
  }
 
void DataIterator_Init(DataIter *self, size_t batch_size, size_t n_batches) {
  self->n = n_batches;
 
  self->lengths = (size_t *)malloc(self->n * sizeof(size_t));
safe_malloc(self->lengths);
  for (size_t i = 0; i < self->n; ++i) {
    self->lengths[i] = batch_size;
  }
 
  self->data = (float **)malloc(self->n * sizeof(float *));
safe_malloc(self->data);
  self->labels = (float **)malloc(self->n * sizeof(float *));
safe_malloc(self->labels);
 
  /* 生成一些随机数据 */
  for (size_t i = 0; i < self->n; ++i) {
    self->data[i] = (float *)malloc(self->lengths[i] * sizeof(float));
safe_malloc(self->data[i]);
    for (size_t j = 0; j < self->lengths[i]; ++j) {
      float x = (float)rand() / (float)(RAND_MAX);
      self->data[i][j] = x;
    }
 
    self->labels[i] = (float *)malloc(self->lengths[i] * sizeof(float));
safe_malloc(self->labels[i]);
    for (size_t j = 0; j < self->lengths[i]; ++j) {
      float y = (float)rand() / (float)(RAND_MAX);
      self->labels[i][j] = y;
    }
  }
 
  self->cur_it = 0;
safe_xgboost(XGProxyDMatrixCreate(&self->_proxy));
}
 
void DataIterator_Free(DataIter *self) {
  for (size_t i = 0; i < self->n; ++i) {
free(self->data[i]);
free(self->labels[i]);
  }
free(self->data);
free(self->lengths);
free(self->labels);
safe_xgboost(XGDMatrixFree(self->_proxy));
};
 
int DataIterator_Next(DataIterHandle handle) {
DataIter *self = (DataIter *)(handle);
  if (self->cur_it == self->n) {
    self->cur_it = 0;
    return 0; /* 到达末尾 */
  }
 
  /* 编码数组接口的 JSON 字符串（来自 numpy 标准） */
  char array[] = "{\"data\": [%lu, false], \"shape\":[%lu, 1], \"typestr\": "
                 "\"<f4\", \"version\": 3}";
memset(self->_array, '\0', sizeof(self->_array));
sprintf(self->_array, array, (size_t)self->data[self->cur_it],
self->lengths[self->cur_it]);
 
safe_xgboost(XGProxyDMatrixSetDataDense(self->_proxy, self->_array));
  /* 传递给迭代器的数据必须保持有效（在下一次迭代或重置之前不被释放） */
* iteration or reset) */
safe_xgboost(XGDMatrixSetDenseInfo(self->_proxy, "label",
self->labels[self->cur_it],
self->lengths[self->cur_it], 1));
  self->cur_it++;
  return 1; /* 继续 */
}
 
void DataIterator_Reset(DataIterHandle handle) {
DataIter *self = (DataIter *)(handle);
  self->cur_it = 0;
}
 
void TrainModel(DMatrix Xy) {
  /* 创建用于训练的 booster */
Booster booster;
DMatrix cache[] = {Xy};
safe_xgboost(XGBoosterCreate(cache, 1, &booster));
  /* 对于外部内存训练，使用 approx 或 hist */
safe_xgboost(XGBoosterSetParam(booster, "tree_method", "hist"));
safe_xgboost(XGBoosterSetParam(booster, "objective", "reg:squarederror"));
 
  /* 开始训练 */
  const char *validation_names[1] = {"train"};
  const char *validation_result = NULL;
  size_t n_rounds = 10;
  for (size_t i = 0; i < n_rounds; ++i) {
safe_xgboost(XGBoosterUpdateOneIter(booster, i, Xy));
safe_xgboost(XGBoosterEvalOneIter(booster, i, cache, validation_names, 1,
&validation_result));
printf("%s\n", validation_result);
  }
 
  /* 将模型保存到 JSON 文件 */
safe_xgboost(XGBoosterSaveModel(booster, "model.json"));
 
safe_xgboost(XGBoosterFree(booster));
}
 
int main() {
DataIter iter;
DataIterator_Init(&iter, BATCH_LEN, N_BATCHS);
 
  /* 从迭代器创建 DMatrix。训练期间，会在当前目录生成一些以 "cache-" 为前缀的缓存文件 */
* prefix "cache-" will be generated in current directory */
  char config[] = "{\"missing\": NaN, \"cache_prefix\": \"cache\"}";
DMatrix Xy;
safe_xgboost(XGDMatrixCreateFromCallback(
&iter, iter._proxy, DataIterator_Reset, DataIterator_Next, config, &Xy));
 
TrainModel(Xy);
 
safe_xgboost(XGDMatrixFree(Xy));
 
DataIterator_Free(&iter);
  return 0;
}