apple / apple/coremltools

Output Discrepancy in CoreML Conversion of XGBoost with objective:'survival:aft'

Open
#2,094 0 comments 0 reactions 0 assignees View on GitHub
bug xgboost
Dominant language
Python
Stars
5.4k
Forks
850
Avg merge
4d 5h
Merged PRs (30d)
10

Description

## 🐞Describing the bug

I am trying to convert an AFT-XGBoost using coremltools.converters.xgboost.convert. When testing the converted model, its predictions have a constant error (or bias) of 1.19314 compared to the predictions of python's model (using output_margin=True).

I have tested several AFT-XGBoost models with different parameters and datasets, having the same problem in all of them. I have tried also a regression XGBoost with objective: 'reg:squarederror', which works completely fine. Thus, I guess it has to be with the differences between the regression model and AFT model.

## To Reproduce

Here is an example to reproduce the issue using the following dataset.
[veterans_lung_cancer.csv](https://github.com/apple/coremltools/files/13422477/veterans_lung_cancer.csv)

```
import numpy as np
import pandas as pd
from sklearn.model_selection import ShuffleSplit

import xgboost as xgb

# The Veterans' Administration Lung Cancer Trial
# The Statistical Analysis of Failure Time Data by Kalbfleisch J. and Prentice R (1980)
df=pd.read_csv("veterans_lung_cancer.csv")
y_lower_bound = df["y_lower_bound"].values
y_upper_bound = df["y_upper_bound"].values
X= df.drop(['y_lower_bound', 'y_upper_bound'], axis=1)

print('Training data:')
print(df)

# Split data into training and validation sets
rs = ShuffleSplit(n_splits=2, test_size=.7, random_state=0)
train_index, valid_index = next(rs.split(X))
dtrain = xgb.DMatrix(X.iloc[train_index, :])
dtrain.set_float_info('label_lower_bound', y_lower_bound[train_index])
dtrain.set_float_info('label_upper_bound', y_upper_bound[train_index])
dvalid = xgb.DMatrix(X.iloc[valid_index, :])
dvalid.set_float_info('label_lower_bound', y_lower_bound[valid_index])
dvalid.set_float_info('label_upper_bound', y_upper_bound[valid_index])

# Train gradient boosted trees using AFT loss and metric
params = {'verbosity': 0,
'objective': 'survival:aft',
'eval_metric': 'aft-nloglik',
'tree_method': 'hist',
'learning_rate': 0.05,
'aft_loss_distribution': 'normal',
'aft_loss_distribution_scale': 1.20,
'max_depth': 6,
'lambda': 0.01,
'alpha': 0.02}
bst = xgb.train(params, dtrain, num_boost_round=10000,
evals=[(dtrain, 'train'), (dvalid, 'valid')],
early_stopping_rounds=50)

from coremltools import converters
import json
coreml_model = converters.xgboost.convert(bst, feature_names=bst.feature_names)
coreml_model.save('aft_model.mlmodel')

Xvalid = X.iloc[valid_index, :].reset_index(drop=True)
pred = np.array(bst.predict(dvalid, output_margin=True))
res_list=[row.to_dict() for i, row in Xvalid.iterrows()]

pred_coreml = np.array([p['target'] for p in coreml_model.predict(res_list)])

df_res = pd.DataFrame({"py-model":pred, "coreml-model":pred_coreml, "diff":pred_coreml-pred})

print(df_res)

df_res.to_csv('res.csv')
```

```
output:

py-model coreml-model diff
0 4.296977 5.490124 1.193148
1 2.702477 3.895624 1.193147
2 6.608140 7.801287 1.193147
3 6.033964 7.227112 1.193148
4 5.036701 6.229849 1.193147
.. ... ... ...
91 3.919354 5.112500 1.193146
92 3.152392 4.345540 1.193148
93 2.697831 3.890977 1.193146
94 4.706688 5.899835 1.193147
95 4.460464 5.653612 1.193148
```

## System environment (please complete the following information):
- coremltools version: 7.1
- xgboost version: 1.4.2
- numpy version: 1.24.4
- OS (e.g. MacOS version or Linux type): MacOS 14.1.1 (I have tried the same experiment with a Ubuntu machine)

Contributor guide

Open the contributing guide

Research direction

Reproduce the discrepancy with the supplied veterans_lung_cancer.csv and the shown converters.xgboost.convert entry point, comparing XGBoost output_margin predictions with the converted model's target output. Trace the AFT conversion path and its handling of survival:aft models. Done means the converted predictions match the Python predictions without the constant 1.19314 bias while reg:squarederror conversion remains unaffected.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.