fix memory size function

Xtra-Computing · Oct 28, 2018 · 235979e · 235979e
1 parent a524e6a
commit 235979e
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/include/thundersvm/model/svmmodel.h b/include/thundersvm/model/svmmodel.h
@@ -90,9 +90,12 @@ class SvmModel {
     //return dec_values
     const SyncArray<float_type> &get_dec_value() const;
 
-    //set max_memory_size during training and prediction
+    //set max_memory_size during training and prediction (input unit MB)
     void set_max_memory_size(size_t size);
 
+	//set max_memory_size during training and prediction (input unit Byte)
+	void set_max_memory_size_Byte(size_t size);
+
     //return n_binary_models
     int get_n_binary_models() const;
 

diff --git a/src/thundersvm/model/svmmodel.cpp b/src/thundersvm/model/svmmodel.cpp
@@ -338,5 +338,10 @@ int SvmModel::get_working_set_size(int n_instances, int n_features) {
 
 void SvmModel::set_max_memory_size(size_t size) {
     if(size > 0)
+		this->param.max_mem_size = static_cast<size_t>(size) << 20;
+}
+
+void SvmModel::set_max_memory_size_Byte(size_t size) {
+	if(size > 0)
 		this->param.max_mem_size = static_cast<size_t>(size);
 }
diff --git a/src/thundersvm/thundersvm-predict.cpp b/src/thundersvm/thundersvm-predict.cpp
@@ -52,7 +52,7 @@ int main(int argc, char **argv) {
         CUDA_CHECK(cudaSetDevice(parser.gpu_id));
 #endif
 
-        model->set_max_memory_size(parser.param_cmd.max_mem_size);
+        model->set_max_memory_size_Byte(parser.param_cmd.max_mem_size);
         model->load_from_file(parser.svmpredict_model_file_name);
         file.close();
         file.open(parser.svmpredict_output_file, fstream::out);