fix prob_predict on cmd #149

Xtra-Computing · May 26, 2019 · 6a62fdd · 6a62fdd
1 parent c7f6882
commit 6a62fdd
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/include/thundersvm/model/svmmodel.h b/include/thundersvm/model/svmmodel.h
@@ -114,6 +114,9 @@ class SvmModel {
 
     //return sv_indices
     const vector<int> &get_sv_ind() const;
+
+    //return label
+    const vector<int> &get_label() const;
 protected:
 
     /**

diff --git a/src/thundersvm/model/svmmodel.cpp b/src/thundersvm/model/svmmodel.cpp
@@ -405,3 +405,7 @@ void SvmModel::compute_linear_coef_single_model(size_t n_feature, const bool zer
 int SvmModel::get_sv_max_index() const{
     return sv_max_index;
 }
+
+const vector<int> &SvmModel::get_label() const{
+    return label;
+}
diff --git a/src/thundersvm/thundersvm-predict.cpp b/src/thundersvm/thundersvm-predict.cpp
@@ -61,8 +61,21 @@ int main(int argc, char **argv) {
 
         vector<float_type> predict_y;
         predict_y = model->predict(predict_dataset.instances(), -1);
-        for (int i = 0; i < predict_y.size(); ++i) {
-            file << predict_y[i] << std::endl;
+        vector<int> label;
+        label = model->get_label();
+        vector<float> prob_predict;
+        prob_predict = model->get_prob_predict();
+        file<<"labels ";
+        for (int i = 0; i < label.size(); i++){
+            file<<label[i]<<" ";
+        }
+        file<<std::endl;
+        for (int i = 0; i < predict_y.size(); i++) {
+            file << predict_y[i]<<" ";
+            for(int j = 0; j < label.size(); j++){
+                file<<prob_predict[label.size() * i + j]<<" ";
+            }
+            file << std::endl;
         }
         file.close();