importargparsefromtypingimportDict,Tupleimportnumpyasnpfrommatplotlibimportpyplotaspltimportxgboostasxgbnp.random.seed(1994)kRows=100kCols=10kClasses=4# number of classeskRounds=10# number of boosting rounds.# Generate some random data for demo.X=np.random.randn(kRows,kCols)y=np.random.randint(0,4,size=kRows)m=xgb.DMatrix(X,y)defsoftmax(x:np.ndarray)->np.ndarray:"""Softmax function with x as input vector."""e=np.exp(x)returne/np.sum(e)defsoftprob_obj(predt:np.ndarray,data:xgb.DMatrix)->Tuple[np.ndarray,np.ndarray]:"""Loss function. Computing the gradient and upper bound on the    Hessian with a diagonal structure for XGBoost (note that this is    not the true Hessian).    Reimplements the `multi:softprob` inside XGBoost.    """labels=data.get_label()ifdata.get_weight().size==0:# Use 1 as weight if we don't have custom weight.weights=np.ones((kRows,1),dtype=float)else:weights=data.get_weight()# The prediction is of shape (rows, classes), each element in a row# represents a raw prediction (leaf weight, hasn't gone through softmax# yet).  In XGBoost 1.0.0, the prediction is transformed by a softmax# function, fixed in later versions.assertpredt.shape==(kRows,kClasses)grad=np.zeros((kRows,kClasses),dtype=float)hess=np.zeros((kRows,kClasses),dtype=float)eps=1e-6# compute the gradient and hessian upper bound, slow iterations in Python, only# suitable for demo.  Also the one in native XGBoost core is more robust to# numeric overflow as we don't do anything to mitigate the `exp` in# `softmax` here.forrinrange(predt.shape[0]):target=labels[r]p=softmax(predt[r,:])forcinrange(predt.shape[1]):asserttarget>=0ortarget<=kClassesg=p[c]-1.0ifc==targetelsep[c]g=g*weights[r]h=max((2.0*p[c]*(1.0-p[c])*weights[r]).item(),eps)grad[r,c]=ghess[r,c]=h# After 2.1.0, pass the gradient as it is.returngrad,hessdefpredict(booster:xgb.Booster,X:xgb.DMatrix)->np.ndarray:"""A customized prediction function that converts raw prediction to    target class.    """# Output margin means we want to obtain the raw prediction obtained from# tree leaf weight.predt=booster.predict(X,output_margin=True)out=np.zeros(kRows)forrinrange(predt.shape[0]):# the class with maximum prob (not strictly prob as it haven't gone# through softmax yet so it doesn't sum to 1, but result is the same# for argmax).i=np.argmax(predt[r])out[r]=ireturnoutdefmerror(predt:np.ndarray,dtrain:xgb.DMatrix)->Tuple[str,np.float64]:y=dtrain.get_label()# Like custom objective, the predt is untransformed leaf weight when custom# objective is provided.# With the use of `custom_metric` parameter in train function, custom metric# receives raw input only when custom objective is also being used.  Otherwise# custom metric will receive transformed prediction.assertpredt.shape==(kRows,kClasses)out=np.zeros(kRows)forrinrange(predt.shape[0]):i=np.argmax(predt[r])out[r]=iasserty.shape==out.shapeerrors=np.zeros(kRows)errors[y!=out]=1.0return"PyMError",np.sum(errors)/kRowsdefplot_history(custom_results:Dict[str,Dict],native_results:Dict[str,Dict])->None:axs:np.ndarrayfig,axs=plt.subplots(2,1)# type: ignoreax0=axs[0]ax1=axs[1]pymerror=custom_results["train"]["PyMError"]merror=native_results["train"]["merror"]x=np.arange(0,kRounds,1)ax0.plot(x,pymerror,label="Custom objective")ax0.legend()ax1.plot(x,merror,label="multi:softmax")ax1.legend()plt.show()defmain(args:argparse.Namespace)->None:# Since 3.1, XGBoost can estimate the base_score automatically for built-in# multi-class objectives.## We explicitly specify it here to disable the automatic estimation to have a proper# comparison between the custom implementation and the built-in implementation.intercept=np.full(shape=(kClasses,),fill_value=1/kClasses)custom_results:Dict[str,Dict]={}# Use our custom objective functionbooster_custom=xgb.train({"num_class":kClasses,"base_score":intercept,"disable_default_eval_metric":True,},m,num_boost_round=kRounds,obj=softprob_obj,custom_metric=merror,evals_result=custom_results,evals=[(m,"train")],)predt_custom=predict(booster_custom,m)native_results:Dict[str,Dict]={}# Use the same objective function defined in XGBoost.booster_native=xgb.train({"num_class":kClasses,"base_score":intercept,"objective":"multi:softmax","eval_metric":"merror",},m,num_boost_round=kRounds,evals_result=native_results,evals=[(m,"train")],)predt_native=booster_native.predict(m)# We are reimplementing the loss function in XGBoost, so it should# be the same for normal cases.assertnp.all(predt_custom==predt_native)np.testing.assert_allclose(custom_results["train"]["PyMError"],native_results["train"]["merror"])ifargs.plot!=0:plot_history(custom_results,native_results)if__name__=="__main__":parser=argparse.ArgumentParser(description="Arguments for custom softmax objective function demo.")parser.add_argument("--plot",type=int,default=1,help="Set to 0 to disable plotting the evaluation history.",)args=parser.parse_args()main(args)

DownloadJupyternotebook:custom_softmax.ipynb

DownloadPythonsourcecode:custom_softmax.py

Downloadzipped:custom_softmax.zip