├── .travis.yml ├── LICENSE ├── Makefile ├── README.md ├── booster.go ├── booster_test.go ├── core ├── README.md ├── utils.go ├── xgboost.go ├── xgboost_test.go ├── xgbooster.go ├── xgdmatrix.go └── xgdmatrix_test.go └── scripts ├── Dockerfile-testing ├── Dockerfile-xgboost ├── build_docker.sh └── docker_run.sh /.travis.yml: -------------------------------------------------------------------------------- 1 | 2 | sudo: required 3 | os: linux 4 | 5 | script: 6 | - make build-docker 7 | - make test -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Unity Technologies Oy 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | build-docker: 2 | ./scripts/build_docker.sh 3 | .PHONY: build-docker 4 | 5 | test: 6 | ./scripts/docker_run.sh go test ./... -bench=. 7 | .PHONY: test -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | WORK IN PROGRESS... USE AT OWN RISK :-) 2 | 3 | [![Build Status](https://travis-ci.org/Applifier/go-xgboost.svg?branch=master)](https://travis-ci.org/Applifier/go-xgboost) 4 | [![GoDoc](https://godoc.org/github.com/Applifier/go-xgboost?status.svg)](http://godoc.org/github.com/Applifier/go-xgboost) 5 | 6 | # go-xgboost 7 | 8 | Go bindings for [XGBoost](https://github.com/dmlc/xgboost) 9 | 10 | ```go 11 | import "github.com/Applifier/go-xgboost" 12 | ``` 13 | 14 | ## Usage 15 | 16 | This library is meant for running predictions against a pre-trained XGBoost model. Limited training related functionality is implemented under [core](https://github.com/Applifier/go-xgboost/blob/master/core) but training the model in python or using the xgboost cli is encouraged. 17 | 18 | ```go 19 | 20 | // Create predictor for a model and define the number of workers (and other settings) 21 | predictor, _ := xgboost.NewPredictor(modelPath, runtime.NumCPU(), 0, 0, -1) 22 | 23 | // Make prediction for one row 24 | res, _ := predictor.Predict(xgboost.FloatSliceVector([]float32{1, 2, 3})) 25 | fmt.Printf("Results: %+v\n", res) 26 | // output: Results: [1.08002] 27 | 28 | ``` 29 | 30 | ## License 31 | 32 | [MIT](https://github.com/Applifier/go-xgboost/blob/master/LICENSE) 33 | -------------------------------------------------------------------------------- /booster.go: -------------------------------------------------------------------------------- 1 | package xgboost 2 | 3 | import ( 4 | "context" 5 | "errors" 6 | "runtime" 7 | 8 | "github.com/Applifier/go-xgboost/core" 9 | ) 10 | 11 | // Matrix interface for 2D matrix 12 | type Matrix interface { 13 | Data() (data []float32, rowCount, columnCount int) 14 | } 15 | 16 | // FloatSliceVector float32 slice backed Matrix implementation 17 | type FloatSliceVector []float32 18 | 19 | // Data returns float32 slice as (1, len(data)) matrix 20 | func (fsm FloatSliceVector) Data() (data []float32, rowCount, columnCount int) { 21 | return fsm, 1, len(fsm) 22 | } 23 | 24 | // Predictor interface for xgboost predictors 25 | type Predictor interface { 26 | Predict(input Matrix) ([]float32, error) 27 | Close(ctx context.Context) error 28 | } 29 | 30 | // NewPredictor returns a new predictor based on given model path, worker count, option mask, ntree_limit and missing value indicator 31 | func NewPredictor(xboostSavedModelPath string, workerCount int, optionMask int, nTreeLimit uint, missingValue float32) (Predictor, error) { 32 | if workerCount <= 0 { 33 | return nil, errors.New("worker count needs to be larger than zero") 34 | } 35 | 36 | requestChan := make(chan multiBoosterRequest) 37 | initErrors := make(chan error) 38 | defer close(initErrors) 39 | 40 | for i := 0; i < workerCount; i++ { 41 | go func() { 42 | runtime.LockOSThread() 43 | defer runtime.UnlockOSThread() 44 | 45 | booster, err := core.XGBoosterCreate(nil) 46 | if err != nil { 47 | initErrors <- err 48 | return 49 | } 50 | 51 | err = booster.LoadModel(xboostSavedModelPath) 52 | if err != nil { 53 | initErrors <- err 54 | return 55 | } 56 | 57 | // No errors occured during init 58 | initErrors <- nil 59 | 60 | for req := range requestChan { 61 | data, rowCount, columnCount := req.matrix.Data() 62 | matrix, err := core.XGDMatrixCreateFromMat(data, rowCount, columnCount, missingValue) 63 | if err != nil { 64 | req.resultChan <- multiBoosterResponse{ 65 | err: err, 66 | } 67 | continue 68 | } 69 | 70 | res, err := booster.Predict(matrix, optionMask, nTreeLimit) 71 | req.resultChan <- multiBoosterResponse{ 72 | err: err, 73 | result: res, 74 | } 75 | } 76 | }() 77 | 78 | err := <-initErrors 79 | if err != nil { 80 | return nil, err 81 | } 82 | } 83 | 84 | return &multiBooster{reqChan: requestChan}, nil 85 | } 86 | 87 | type multiBoosterRequest struct { 88 | matrix Matrix 89 | resultChan chan multiBoosterResponse 90 | } 91 | 92 | type multiBoosterResponse struct { 93 | err error 94 | result []float32 95 | } 96 | 97 | type multiBooster struct { 98 | reqChan chan multiBoosterRequest 99 | } 100 | 101 | func (mb *multiBooster) Predict(input Matrix) ([]float32, error) { 102 | resChan := make(chan multiBoosterResponse) 103 | mb.reqChan <- multiBoosterRequest{ 104 | matrix: input, 105 | resultChan: resChan, 106 | } 107 | 108 | result := <-resChan 109 | return result.result, result.err 110 | } 111 | 112 | func (mb *multiBooster) Close(ctx context.Context) error { 113 | close(mb.reqChan) 114 | return nil 115 | } 116 | -------------------------------------------------------------------------------- /booster_test.go: -------------------------------------------------------------------------------- 1 | package xgboost 2 | 3 | import ( 4 | "context" 5 | "fmt" 6 | "io/ioutil" 7 | "log" 8 | "os" 9 | "path" 10 | "runtime" 11 | "testing" 12 | 13 | "github.com/Applifier/go-xgboost/core" 14 | ) 15 | 16 | type tester interface { 17 | Helper() 18 | Error(args ...interface{}) 19 | } 20 | 21 | func trainAndSaveModel(t tester) (string, func()) { 22 | if t != nil { 23 | t.Helper() 24 | } 25 | cols := 3 26 | rows := 5 27 | trainData := make([]float32, cols*rows) 28 | for i := 0; i < rows; i++ { 29 | for j := 0; j < cols; j++ { 30 | trainData[(i*cols)+j] = float32((i + 1) * (j + 1)) 31 | } 32 | } 33 | 34 | trainLabels := make([]float32, rows) 35 | for i := 0; i < rows; i++ { 36 | trainLabels[i] = float32(1 + i*i*i) 37 | } 38 | 39 | matrix, err := core.XGDMatrixCreateFromMat(trainData, rows, cols, -1) 40 | if err != nil && t != nil { 41 | t.Error(err) 42 | } 43 | 44 | err = matrix.SetFloatInfo("label", trainLabels) 45 | if err != nil && t != nil { 46 | t.Error(err) 47 | } 48 | 49 | booster, err := core.XGBoosterCreate([]*core.XGDMatrix{matrix}) 50 | if err != nil && t != nil { 51 | t.Error(err) 52 | } 53 | 54 | noErr := func(err error) { 55 | if err != nil && t != nil { 56 | t.Error(err) 57 | } 58 | } 59 | 60 | noErr(booster.SetParam("booster", "gbtree")) 61 | noErr(booster.SetParam("objective", "reg:linear")) 62 | noErr(booster.SetParam("max_depth", "5")) 63 | noErr(booster.SetParam("eta", "0.1")) 64 | noErr(booster.SetParam("min_child_weight", "1")) 65 | noErr(booster.SetParam("subsample", "0.5")) 66 | noErr(booster.SetParam("colsample_bytree", "1")) 67 | noErr(booster.SetParam("num_parallel_tree", "1")) 68 | noErr(booster.SetParam("silent", "1")) 69 | 70 | // perform 200 learning iterations 71 | for iter := 0; iter < 200; iter++ { 72 | noErr(booster.UpdateOneIter(iter, matrix)) 73 | } 74 | 75 | dir, err := ioutil.TempDir("", "go-xgboost") 76 | if err != nil { 77 | log.Fatal(err) 78 | } 79 | 80 | savePath := path.Join(dir, "testmodel.bst") 81 | 82 | noErr(booster.SaveModel(savePath)) 83 | 84 | return savePath, func() { 85 | os.RemoveAll(dir) 86 | } 87 | } 88 | 89 | func TestBooster(t *testing.T) { 90 | modelPath, cleanUp := trainAndSaveModel(t) 91 | defer cleanUp() 92 | 93 | predictor, err := NewPredictor(modelPath, 1, 0, 0, -1) 94 | if err != nil { 95 | t.Fatal(err) 96 | } 97 | defer predictor.Close(context.TODO()) 98 | 99 | cols := 3 100 | rows := 5 101 | testData := make([][]float32, rows) 102 | for i := 0; i < rows; i++ { 103 | testData[i] = make([]float32, cols) 104 | for j := 0; j < cols; j++ { 105 | testData[i][j] = float32((i + 1) * (j + 1)) 106 | } 107 | } 108 | 109 | expectedResult := []float32{1.0631807, 2.4375393, 8.3054695, 30.843433, 63.097855} 110 | 111 | for i, test := range testData { 112 | res, err := predictor.Predict(FloatSliceVector(test)) 113 | if err != nil { 114 | t.Error(err) 115 | } 116 | 117 | if res[0] != expectedResult[i] { 118 | t.Error("unexpected result received", expectedResult[i]) 119 | } 120 | } 121 | } 122 | 123 | func BenchmarkBooster(b *testing.B) { 124 | modelPath, cleanUp := trainAndSaveModel(b) 125 | defer cleanUp() 126 | 127 | predictor, err := NewPredictor(modelPath, 1, 0, 0, -1) 128 | if err != nil { 129 | b.Fatal(err) 130 | } 131 | defer predictor.Close(context.TODO()) 132 | 133 | testData := []float32{1, 2, 3} 134 | 135 | b.ResetTimer() 136 | 137 | for i := 0; i < b.N; i++ { 138 | res, err := predictor.Predict(FloatSliceVector(testData)) 139 | if err != nil { 140 | b.Error(err) 141 | } 142 | if len(res) != 1 { 143 | b.Error("invalid amount of results received") 144 | } 145 | } 146 | } 147 | 148 | func BenchmarkBoosterParallel(b *testing.B) { 149 | modelPath, cleanUp := trainAndSaveModel(b) 150 | defer cleanUp() 151 | 152 | predictor, err := NewPredictor(modelPath, runtime.NumCPU(), 0, 0, -1) 153 | if err != nil { 154 | b.Fatal(err) 155 | } 156 | defer predictor.Close(context.TODO()) 157 | 158 | testData := []float32{1, 2, 3} 159 | 160 | b.ResetTimer() 161 | 162 | b.RunParallel(func(pb *testing.PB) { 163 | for pb.Next() { 164 | res, err := predictor.Predict(FloatSliceVector(testData)) 165 | if err != nil { 166 | b.Error(err) 167 | } 168 | if len(res) != 1 { 169 | b.Error("invalid amount of results received") 170 | } 171 | } 172 | }) 173 | } 174 | 175 | func ExampleBooster() { 176 | // Retrieve filepath for a pre-trained model 177 | modelPath, cleanUp := trainAndSaveModel(nil) 178 | defer cleanUp() 179 | 180 | // Create predictor and define the number of workers (and other settings) 181 | predictor, _ := NewPredictor(modelPath, runtime.NumCPU(), 0, 0, -1) 182 | defer predictor.Close(context.TODO()) 183 | 184 | res, _ := predictor.Predict(FloatSliceVector([]float32{1, 2, 3})) 185 | fmt.Printf("Results: %+v\n", res) 186 | // output: Results: [1.0631807] 187 | } 188 | -------------------------------------------------------------------------------- /core/README.md: -------------------------------------------------------------------------------- 1 | [![GoDoc](https://godoc.org/github.com/Applifier/go-xgboost/core?status.svg)](http://godoc.org/github.com/Applifier/go-xgboost/core) 2 | 3 | # Core package 4 | 5 | ```go 6 | import "github.com/Applifier/go-xgboost/core" 7 | ``` 8 | 9 | ## Example 10 | 11 | ```go 12 | // create the training data 13 | cols := 3 14 | rows := 5 15 | trainData := make([]float32, cols*rows) 16 | for i := 0; i < rows; i++ { 17 | for j := 0; j < cols; j++ { 18 | trainData[(i*cols)+j] = float32((i + 1) * (j + 1)) 19 | } 20 | } 21 | 22 | trainLabels := make([]float32, rows) 23 | for i := 0; i < rows; i++ { 24 | trainLabels[i] = float32(1 + i*i*i) 25 | } 26 | 27 | // Create XGDMatrix for training data 28 | matrix, _ := core.XGDMatrixCreateFromMat(trainData, rows, cols, -1) 29 | 30 | // Set training labels 31 | matrix.SetFloatInfo("label", trainLabels) 32 | 33 | // Create booster 34 | booster, _ := core.XGBoosterCreate([]*core.XGDMatrix{matrix}) 35 | 36 | // Set booster parameters 37 | booster.SetParam("booster", "gbtree") 38 | booster.SetParam("objective", "reg:linear") 39 | booster.SetParam("max_depth", "5") 40 | booster.SetParam("eta", "0.1") 41 | booster.SetParam("min_child_weight", "1") 42 | booster.SetParam("subsample", "0.5") 43 | booster.SetParam("colsample_bytree", "1") 44 | booster.SetParam("num_parallel_tree", "1") 45 | 46 | // perform 200 learning iterations 47 | for iter := 0; iter < 200; iter++ { 48 | booster.UpdateOneIter(iter, matrix) 49 | } 50 | 51 | testData := make([]float32, cols*rows) 52 | for i := 0; i < rows; i++ { 53 | for j := 0; j < cols; j++ { 54 | testData[(i*cols)+j] = float32((i + 1) * (j + 1)) 55 | } 56 | } 57 | 58 | // Create XGDMatrix for test data 59 | testmat, _ := core.XGDMatrixCreateFromMat(testData, rows, cols, -1) 60 | 61 | // Predict 62 | res, _ := booster.Predict(testmat, 0, 0) 63 | 64 | fmt.Printf("%+v\n", res) 65 | // output: [1.08002 2.5686886 7.86032 29.923136 63.76062] 66 | ``` -------------------------------------------------------------------------------- /core/utils.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | func copyUint32Slice(sli []uint32) []uint32 { 4 | n := make([]uint32, len(sli)) 5 | copy(n, sli) 6 | return n 7 | } 8 | 9 | func copyFloat32Slice(sli []float32) []float32 { 10 | n := make([]float32, len(sli)) 11 | copy(n, sli) 12 | return n 13 | } 14 | -------------------------------------------------------------------------------- /core/xgboost.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | /* 4 | #cgo LDFLAGS: -lxgboost 5 | #include 6 | */ 7 | import "C" 8 | import ( 9 | "errors" 10 | ) 11 | 12 | func checkError(res C.int) error { 13 | if int(res) != 0 { 14 | errStr := C.XGBGetLastError() 15 | return errors.New(C.GoString(errStr)) 16 | } 17 | 18 | return nil 19 | } -------------------------------------------------------------------------------- /core/xgboost_test.go: -------------------------------------------------------------------------------- 1 | package core_test 2 | 3 | import ( 4 | "fmt" 5 | "io/ioutil" 6 | "log" 7 | "math" 8 | "os" 9 | "path" 10 | "testing" 11 | 12 | "github.com/Applifier/go-xgboost/core" 13 | ) 14 | 15 | func TestXGBoost(t *testing.T) { 16 | // create the training data 17 | cols := 3 18 | rows := 5 19 | trainData := make([]float32, cols*rows) 20 | for i := 0; i < rows; i++ { 21 | for j := 0; j < cols; j++ { 22 | trainData[(i*cols)+j] = float32((i + 1) * (j + 1)) 23 | } 24 | } 25 | 26 | trainLabels := make([]float32, rows) 27 | for i := 0; i < rows; i++ { 28 | trainLabels[i] = float32(1 + i*i*i) 29 | } 30 | 31 | matrix, err := core.XGDMatrixCreateFromMat(trainData, rows, cols, -1) 32 | if err != nil { 33 | t.Error(err) 34 | } 35 | 36 | err = matrix.SetFloatInfo("label", trainLabels) 37 | if err != nil { 38 | t.Error(err) 39 | } 40 | 41 | booster, err := core.XGBoosterCreate([]*core.XGDMatrix{matrix}) 42 | if err != nil { 43 | t.Error(err) 44 | } 45 | 46 | noErr := func(err error) { 47 | if err != nil { 48 | t.Error(err) 49 | } 50 | } 51 | 52 | noErr(booster.SetParam("booster", "gbtree")) 53 | noErr(booster.SetParam("objective", "reg:linear")) 54 | noErr(booster.SetParam("max_depth", "5")) 55 | noErr(booster.SetParam("eta", "0.1")) 56 | noErr(booster.SetParam("min_child_weight", "1")) 57 | noErr(booster.SetParam("subsample", "0.5")) 58 | noErr(booster.SetParam("colsample_bytree", "1")) 59 | noErr(booster.SetParam("num_parallel_tree", "1")) 60 | noErr(booster.SetParam("silent", "1")) 61 | 62 | // perform 200 learning iterations 63 | for iter := 0; iter < 200; iter++ { 64 | noErr(booster.UpdateOneIter(iter, matrix)) 65 | } 66 | 67 | testData := make([]float32, cols*rows) 68 | for i := 0; i < rows; i++ { 69 | for j := 0; j < cols; j++ { 70 | testData[(i*cols)+j] = float32((i + 1) * (j + 1)) 71 | } 72 | } 73 | 74 | testmat, err := core.XGDMatrixCreateFromMat(testData, rows, cols, -1) 75 | if err != nil { 76 | t.Error(err) 77 | } 78 | 79 | res, err := booster.Predict(testmat, 0, 0) 80 | if err != nil { 81 | t.Error(err) 82 | } 83 | 84 | // TODO measure actual accuracy 85 | totalDiff := 0.0 86 | for i, label := range trainLabels { 87 | diff := math.Abs(float64(label - res[i])) 88 | totalDiff += diff 89 | } 90 | 91 | if totalDiff > 6.0 { 92 | t.Error("error is too large") 93 | } 94 | 95 | dir, err := ioutil.TempDir("", "go-xgboost") 96 | if err != nil { 97 | log.Fatal(err) 98 | } 99 | defer os.RemoveAll(dir) // clean up 100 | 101 | savePath := path.Join(dir, "testmodel.bst") 102 | 103 | noErr(booster.SaveModel(savePath)) 104 | 105 | newBooster, err := core.XGBoosterCreate(nil) 106 | if err != nil { 107 | t.Error(err) 108 | } 109 | 110 | noErr(newBooster.LoadModel(savePath)) 111 | 112 | testmat2, err := core.XGDMatrixCreateFromMat(testData, rows, cols, -1) 113 | if err != nil { 114 | t.Error(err) 115 | } 116 | 117 | res, err = newBooster.Predict(testmat2, 0, 0) 118 | if err != nil { 119 | t.Error(err) 120 | } 121 | 122 | // TODO measure actual accuracy 123 | totalDiff = 0.0 124 | for i, label := range trainLabels { 125 | diff := math.Abs(float64(label - res[i])) 126 | totalDiff += diff 127 | } 128 | 129 | if totalDiff > 6.0 { 130 | t.Error("error is too large") 131 | } 132 | } 133 | 134 | func ExampleXGBoost() { 135 | // create the train data 136 | cols := 3 137 | rows := 5 138 | trainData := make([]float32, cols*rows) 139 | for i := 0; i < rows; i++ { 140 | for j := 0; j < cols; j++ { 141 | trainData[(i*cols)+j] = float32((i + 1) * (j + 1)) 142 | } 143 | } 144 | 145 | trainLabels := make([]float32, rows) 146 | for i := 0; i < rows; i++ { 147 | trainLabels[i] = float32(1 + i*i*i) 148 | } 149 | 150 | // Create XGDMatrix for training data 151 | matrix, _ := core.XGDMatrixCreateFromMat(trainData, rows, cols, -1) 152 | 153 | // Set training labels 154 | matrix.SetFloatInfo("label", trainLabels) 155 | 156 | // Create booster 157 | booster, _ := core.XGBoosterCreate([]*core.XGDMatrix{matrix}) 158 | 159 | // Set booster parameters 160 | booster.SetParam("booster", "gbtree") 161 | booster.SetParam("objective", "reg:linear") 162 | booster.SetParam("max_depth", "5") 163 | booster.SetParam("eta", "0.1") 164 | booster.SetParam("min_child_weight", "1") 165 | booster.SetParam("subsample", "0.5") 166 | booster.SetParam("colsample_bytree", "1") 167 | booster.SetParam("num_parallel_tree", "1") 168 | booster.SetParam("silent", "1") 169 | 170 | // perform 200 learning iterations 171 | for iter := 0; iter < 200; iter++ { 172 | booster.UpdateOneIter(iter, matrix) 173 | } 174 | 175 | testData := make([]float32, cols*rows) 176 | for i := 0; i < rows; i++ { 177 | for j := 0; j < cols; j++ { 178 | testData[(i*cols)+j] = float32((i + 1) * (j + 1)) 179 | } 180 | } 181 | 182 | // Create XGDMatrix for test data 183 | testmat, _ := core.XGDMatrixCreateFromMat(testData, rows, cols, -1) 184 | 185 | // Predict 186 | res, _ := booster.Predict(testmat, 0, 0) 187 | 188 | fmt.Printf("%+v\n", res) 189 | // output: [1.0631807 2.4375393 8.3054695 30.843433 63.097855] 190 | } 191 | -------------------------------------------------------------------------------- /core/xgbooster.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | /* 4 | #cgo LDFLAGS: -lxgboost 5 | #include 6 | #include 7 | #include 8 | */ 9 | import "C" 10 | import ( 11 | "reflect" 12 | "runtime" 13 | "unsafe" 14 | ) 15 | 16 | // XGBooster gradient booster 17 | type XGBooster struct { 18 | handle C.BoosterHandle 19 | } 20 | 21 | // SetParam set parameters 22 | func (booster *XGBooster) SetParam(name string, value string) error { 23 | cname := C.CString(name) 24 | defer C.free(unsafe.Pointer(cname)) 25 | 26 | cvalue := C.CString(value) 27 | defer C.free(unsafe.Pointer(cvalue)) 28 | 29 | res := C.XGBoosterSetParam(booster.handle, cname, cvalue) 30 | if err := checkError(res); err != nil { 31 | return err 32 | } 33 | 34 | return nil 35 | } 36 | 37 | // DeleteParam set parameters 38 | func (booster *XGBooster) DeleteParam(name string) error { 39 | cname := C.CString(name) 40 | defer C.free(unsafe.Pointer(cname)) 41 | 42 | res := C.XGBoosterSetParam(booster.handle, cname, nil) 43 | if err := checkError(res); err != nil { 44 | return err 45 | } 46 | 47 | return nil 48 | } 49 | 50 | // UpdateOneIter update the model in one round using dtrain 51 | func (booster *XGBooster) UpdateOneIter(iter int, mat *XGDMatrix) error { 52 | res := C.XGBoosterUpdateOneIter(booster.handle, C.int(iter), mat.handle) 53 | if err := checkError(res); err != nil { 54 | return err 55 | } 56 | 57 | return nil 58 | } 59 | 60 | // Predict make prediction based on dmat 61 | func (booster *XGBooster) Predict(mat *XGDMatrix, optionMask int, ntreeLimit uint) ([]float32, error) { 62 | var outLen C.bst_ulong 63 | var outResult *C.float 64 | 65 | res := C.XGBoosterPredict(booster.handle, mat.handle, C.int(optionMask), C.uint(ntreeLimit), &outLen, &outResult) 66 | if err := checkError(res); err != nil { 67 | return nil, err 68 | } 69 | 70 | var list []float32 71 | sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&list))) 72 | sliceHeader.Cap = int(outLen) 73 | sliceHeader.Len = int(outLen) 74 | sliceHeader.Data = uintptr(unsafe.Pointer(outResult)) 75 | 76 | runtime.KeepAlive(mat) 77 | 78 | return copyFloat32Slice(list), nil 79 | } 80 | 81 | // LoadModel load model from existing file 82 | func (booster *XGBooster) LoadModel(filePath string) error { 83 | cfilePath := C.CString(filePath) 84 | defer C.free(unsafe.Pointer(cfilePath)) 85 | 86 | return checkError(C.XGBoosterLoadModel(booster.handle, cfilePath)) 87 | } 88 | 89 | // SaveModel save model into file 90 | func (booster *XGBooster) SaveModel(filePath string) error { 91 | cfilePath := C.CString(filePath) 92 | defer C.free(unsafe.Pointer(cfilePath)) 93 | 94 | return checkError(C.XGBoosterSaveModel(booster.handle, cfilePath)) 95 | } 96 | 97 | func xdgBoosterFinalizer(booster *XGBooster) { 98 | C.XGBoosterFree(booster.handle) 99 | } 100 | 101 | // XGBoosterCreate creates a new booster for a given matrixes 102 | func XGBoosterCreate(matrix []*XGDMatrix) (*XGBooster, error) { 103 | var ptr *C.DMatrixHandle 104 | handles := make([]C.DMatrixHandle, len(matrix)) 105 | for i, matrix := range matrix { 106 | handles[i] = matrix.handle 107 | } 108 | if len(handles) > 0 { 109 | ptr = (*C.DMatrixHandle)(&handles[0]) 110 | } 111 | 112 | var out C.BoosterHandle 113 | res := C.XGBoosterCreate(ptr, C.bst_ulong(len(handles)), &out) 114 | if err := checkError(res); err != nil { 115 | return nil, err 116 | } 117 | 118 | booster := &XGBooster{handle: out} 119 | runtime.SetFinalizer(booster, xdgBoosterFinalizer) 120 | 121 | return booster, nil 122 | } 123 | -------------------------------------------------------------------------------- /core/xgdmatrix.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | /* 4 | #cgo LDFLAGS: -lxgboost 5 | #include 6 | #include 7 | #include 8 | */ 9 | import "C" 10 | import ( 11 | "errors" 12 | "reflect" 13 | "runtime" 14 | "unsafe" 15 | ) 16 | 17 | // ErrNotSuccessful returned when an action fails 18 | var ErrNotSuccessful = errors.New("not succesfull") 19 | 20 | // XGDMatrix matrix 21 | type XGDMatrix struct { 22 | handle C.DMatrixHandle 23 | cols int 24 | rows int 25 | } 26 | 27 | // NumRow get number of rows. 28 | func (matrix *XGDMatrix) NumRow() (uint32, error) { 29 | var count C.bst_ulong 30 | if err := checkError(C.XGDMatrixNumRow(matrix.handle, &count)); err != nil { 31 | return 0, err 32 | } 33 | 34 | return uint32(count), nil 35 | } 36 | 37 | // NumCol get number of cols. 38 | func (matrix *XGDMatrix) NumCol() (uint32, error) { 39 | var count C.bst_ulong 40 | if err := checkError(C.XGDMatrixNumCol(matrix.handle, &count)); err != nil { 41 | return 0, err 42 | } 43 | 44 | return uint32(count), nil 45 | } 46 | 47 | // SetUIntInfo set uint32 vector to a content in info 48 | func (matrix *XGDMatrix) SetUIntInfo(field string, values []uint32) error { 49 | cstr := C.CString(field) 50 | defer C.free(unsafe.Pointer(cstr)) 51 | 52 | res := C.XGDMatrixSetUIntInfo(matrix.handle, cstr, (*C.uint)(&values[0]), C.bst_ulong(len(values))) 53 | 54 | runtime.KeepAlive(values) 55 | return checkError(res) 56 | } 57 | 58 | // SetGroup set label of the training matrix 59 | func (matrix *XGDMatrix) SetGroup(group ...uint32) error { 60 | res := C.XGDMatrixSetGroup(matrix.handle, (*C.uint)(&group[0]), C.bst_ulong(len(group))) 61 | runtime.KeepAlive(group) 62 | return checkError(res) 63 | } 64 | 65 | // SetFloatInfo set float vector to a content in info 66 | func (matrix *XGDMatrix) SetFloatInfo(field string, values []float32) error { 67 | cstr := C.CString(field) 68 | defer C.free(unsafe.Pointer(cstr)) 69 | 70 | res := C.XGDMatrixSetFloatInfo(matrix.handle, cstr, (*C.float)(&values[0]), C.bst_ulong(len(values))) 71 | if err := checkError(res); err != nil { 72 | return err 73 | } 74 | runtime.KeepAlive(values) 75 | 76 | return nil 77 | } 78 | 79 | // GetFloatInfo get float info vector from matrix 80 | func (matrix *XGDMatrix) GetFloatInfo(field string) ([]float32, error) { 81 | cstr := C.CString(field) 82 | defer C.free(unsafe.Pointer(cstr)) 83 | 84 | var outLen C.bst_ulong 85 | var outResult *C.float 86 | 87 | if err := checkError(C.XGDMatrixGetFloatInfo(matrix.handle, cstr, &outLen, &outResult)); err != nil { 88 | return nil, err 89 | } 90 | 91 | var list []float32 92 | sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&list))) 93 | sliceHeader.Cap = int(outLen) 94 | sliceHeader.Len = int(outLen) 95 | sliceHeader.Data = uintptr(unsafe.Pointer(outResult)) 96 | 97 | return copyFloat32Slice(list), nil 98 | } 99 | 100 | // GetUIntInfo get uint32 info vector from matrix 101 | func (matrix *XGDMatrix) GetUIntInfo(field string) ([]uint32, error) { 102 | cstr := C.CString(field) 103 | defer C.free(unsafe.Pointer(cstr)) 104 | 105 | var outLen C.bst_ulong 106 | var outResult *C.uint 107 | 108 | if err := checkError(C.XGDMatrixGetUIntInfo(matrix.handle, cstr, &outLen, &outResult)); err != nil { 109 | return nil, err 110 | } 111 | 112 | var list []uint32 113 | sliceHeader := (*reflect.SliceHeader)((unsafe.Pointer(&list))) 114 | sliceHeader.Cap = int(outLen) 115 | sliceHeader.Len = int(outLen) 116 | sliceHeader.Data = uintptr(unsafe.Pointer(outResult)) 117 | 118 | return copyUint32Slice(list), nil 119 | } 120 | 121 | func xdgMatrixFinalizer(mat *XGDMatrix) { 122 | C.XGDMatrixFree(mat.handle) 123 | } 124 | 125 | // XGDMatrixCreateFromMat create matrix content from dense matrix 126 | func XGDMatrixCreateFromMat(data []float32, nrows int, ncols int, missing float32) (*XGDMatrix, error) { 127 | if len(data) != nrows*ncols { 128 | return nil, errors.New("data length doesn't match given dimensions") 129 | } 130 | 131 | var out C.DMatrixHandle 132 | res := C.XGDMatrixCreateFromMat((*C.float)(&data[0]), C.bst_ulong(nrows), C.bst_ulong(ncols), C.float(missing), &out) 133 | if err := checkError(res); err != nil { 134 | return nil, err 135 | } 136 | 137 | matrix := &XGDMatrix{handle: out, rows: nrows, cols: ncols} 138 | runtime.SetFinalizer(matrix, xdgMatrixFinalizer) 139 | 140 | runtime.KeepAlive(data) 141 | 142 | return matrix, nil 143 | } 144 | -------------------------------------------------------------------------------- /core/xgdmatrix_test.go: -------------------------------------------------------------------------------- 1 | package core 2 | 3 | import ( 4 | "testing" 5 | ) 6 | 7 | func TestXGDMatrix(t *testing.T) { 8 | data := []float32{1, 2, 3, 4} 9 | 10 | matrix, err := XGDMatrixCreateFromMat(data, 2, 2, -1) 11 | if err != nil { 12 | t.Error(err) 13 | } 14 | 15 | if matrix == nil { 16 | t.Error("matrix was not created") 17 | } 18 | 19 | err = matrix.SetFloatInfo("label", []float32{1, 2}) 20 | if err != nil { 21 | t.Error(err) 22 | } 23 | 24 | vals, err := matrix.GetFloatInfo("label") 25 | if err != nil { 26 | t.Error(err) 27 | } 28 | 29 | if vals[0] != 1 || vals[1] != 2 { 30 | t.Error("Wrong values returned") 31 | } 32 | 33 | rowCount, err := matrix.NumRow() 34 | if err != nil { 35 | t.Error(err) 36 | } 37 | 38 | if rowCount != 2 { 39 | t.Error("Wrong row count returned") 40 | } 41 | 42 | colCount, err := matrix.NumCol() 43 | if err != nil { 44 | t.Error(err) 45 | } 46 | 47 | if colCount != 2 { 48 | t.Error("Wrong col count returned") 49 | } 50 | } 51 | -------------------------------------------------------------------------------- /scripts/Dockerfile-testing: -------------------------------------------------------------------------------- 1 | FROM xgboost:latest 2 | 3 | # Install golang 4 | RUN apt-get update 5 | RUN apt-get install -y wget pkg-config git gcc 6 | 7 | RUN wget -P /tmp https://storage.googleapis.com/golang/go1.10.2.linux-amd64.tar.gz 8 | 9 | RUN tar -C /usr/local -xzf /tmp/go1.10.2.linux-amd64.tar.gz 10 | RUN rm /tmp/go1.10.2.linux-amd64.tar.gz 11 | 12 | ENV GOPATH /go 13 | ENV PATH $GOPATH/bin:/usr/local/go/bin:$PATH 14 | RUN mkdir -p "$GOPATH/src" "$GOPATH/bin" && chmod -R 777 "$GOPATH" 15 | 16 | -------------------------------------------------------------------------------- /scripts/Dockerfile-xgboost: -------------------------------------------------------------------------------- 1 | FROM ubuntu:18.04 2 | 3 | # Install essential dependencies 4 | 5 | RUN apt-get update && apt-get install -y \ 6 | build-essential \ 7 | curl \ 8 | libcurl3-dev \ 9 | git \ 10 | libssl-dev && \ 11 | apt-get clean && \ 12 | rm -rf /var/lib/apt/lists/* 13 | 14 | RUN mkdir /src && \ 15 | cd /src && \ 16 | git clone --depth 1 --recursive https://github.com/dmlc/xgboost.git && \ 17 | cd /src/xgboost && make && \ 18 | cp /src/xgboost/lib/libxgboost.so /usr/local/lib && ldconfig -n -v /usr/local/lib && \ 19 | cp -r /src/xgboost/include/xgboost /usr/local/include/xgboost && \ 20 | cp -r /src/xgboost/rabit/include/rabit /usr/local/include/rabit && \ 21 | cp -r /src/xgboost/dmlc-core/include/dmlc /usr/local/include/dmlc && \ 22 | rm -rf /src/xgboost 23 | -------------------------------------------------------------------------------- /scripts/build_docker.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null && pwd )" 4 | 5 | docker build -t xgboost:latest - < "$DIR/Dockerfile-xgboost" 6 | docker build -t xgboost-testing:latest - < "$DIR/Dockerfile-testing" 7 | -------------------------------------------------------------------------------- /scripts/docker_run.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | docker run --rm -it -v "$PWD":/go/src/github.com/Applifier/go-xgboost -w /go/src/github.com/Applifier/go-xgboost xgboost-testing:latest $@ --------------------------------------------------------------------------------