fix predict output on cmd

Xtra-Computing · May 27, 2019 · 70b64f2 · 70b64f2
1 parent 6a62fdd
commit 70b64f2
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 13 deletions.
diff --git a/include/thundersvm/model/svmmodel.h b/include/thundersvm/model/svmmodel.h
@@ -117,6 +117,9 @@ class SvmModel {
 
     //return label
     const vector<int> &get_label() const;
+
+    //return param.probability
+    const bool is_prob() const;
 protected:
 
     /**

diff --git a/src/thundersvm/model/svmmodel.cpp b/src/thundersvm/model/svmmodel.cpp
@@ -339,10 +339,15 @@ const SyncArray<float_type> &SvmModel::get_dec_value() const {
 int SvmModel::get_n_binary_models() const{
     return n_binary_models;
 }
+
 const vector<float> &SvmModel::get_prob_predict() const{
     return prob_predict;
 }
 
+const bool SvmModel::is_prob() const{
+    return param.probability;
+}
+
 int SvmModel::get_working_set_size(int n_instances, int n_features) {
     size_t free_mem = param.max_mem_size - SyncMem::get_total_memory_size();
     int ws_size = min(max2power(n_instances),

diff --git a/src/thundersvm/thundersvm-predict.cpp b/src/thundersvm/thundersvm-predict.cpp
@@ -61,21 +61,29 @@ int main(int argc, char **argv) {
 
         vector<float_type> predict_y;
         predict_y = model->predict(predict_dataset.instances(), -1);
-        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]<<" ";
+
+        if(model->is_prob() == 1) {
+            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;
+            }
+        }
+        else{
+            for (int i = 0; i < predict_y.size(); ++i) {
+                file << predict_y[i] << std::endl;
+            }
         }
         file.close();