option to run on cpu without gpu

diff --git a/include/caffe/common.hpp b/include/caffe/common.hpp
index 7362b233780d7d205906df1e53bbab2950e86fa5..485a86947f5f9834b056593dd60ca6507b3f935e 100644 (file)
--- a/include/caffe/common.hpp
+++ b/include/caffe/common.hpp
@@ -85,7 +85,10 @@ class Caffe {
   // Returns the phase: TRAIN or TEST.
   inline static Phase phase() { return Get().phase_; }
   // The setters for the variables
-  // Sets the mode.
+  // Sets the mode. It is recommended that you don't change the mode halfway
+  // into the program since that may cause allocation of pinned memory being
+  // freed in a non-pinned way, which may cause problems - I haven't verified
+  // it personally but better to note it here in the header file.
   inline static void set_mode(Brew mode) { Get().mode_ = mode; }
   // Sets the phase.
   inline static void set_phase(Phase phase) { Get().phase_ = phase; }

diff --git a/include/caffe/syncedmem.hpp b/include/caffe/syncedmem.hpp
index 862512f9250ca55a63c873ef33e02fc4c796014b..1229adb47759fbfdaab70514b928d1459e59de11 100644 (file)
--- a/include/caffe/syncedmem.hpp
+++ b/include/caffe/syncedmem.hpp
@@ -3,8 +3,30 @@
 #ifndef CAFFE_SYNCEDMEM_HPP_
 #define CAFFE_SYNCEDMEM_HPP_
 
+#include <cstdlib>
+
+#include "caffe/common.hpp"
+
 namespace caffe {
 
+inline void CaffeMallocHost(void** ptr, size_t size) {
+  if (Caffe::mode() == Caffe::GPU) {
+    CUDA_CHECK(cudaMallocHost(ptr, size));
+  } else {
+    *ptr = malloc(size);
+  }
+}
+
+
+inline void CaffeFreeHost(void* ptr) {
+  if (Caffe::mode() == Caffe::GPU) {
+    CUDA_CHECK(cudaFreeHost(ptr)));
+  } else {
+    free(ptr)
+  }
+}
+
+
 class SyncedMemory {
  public:
   SyncedMemory()

diff --git a/src/caffe/syncedmem.cpp b/src/caffe/syncedmem.cpp
index e6600f2119de8b790e724575b088f90f0f31cf85..8f6ecf62a0861b206c34446b61cf37e5fd762e91 100644 (file)
--- a/src/caffe/syncedmem.cpp
+++ b/src/caffe/syncedmem.cpp
@@ -11,7 +11,7 @@ namespace caffe {
 
 SyncedMemory::~SyncedMemory() {
   if (cpu_ptr_) {
-    CUDA_CHECK(cudaFreeHost(cpu_ptr_));
+    CaffeFreeHost(cpu_ptr_);
   }
 
   if (gpu_ptr_) {
@@ -22,13 +22,13 @@ SyncedMemory::~SyncedMemory() {
 inline void SyncedMemory::to_cpu() {
   switch (head_) {
   case UNINITIALIZED:
-    CUDA_CHECK(cudaMallocHost(&cpu_ptr_, size_));
+    CaffeMallocHost(&cpu_ptr_, size_);
     memset(cpu_ptr_, 0, size_);
     head_ = HEAD_AT_CPU;
     break;
   case HEAD_AT_GPU:
     if (cpu_ptr_ == NULL) {
-      CUDA_CHECK(cudaMallocHost(&cpu_ptr_, size_));
+      CaffeMallocHost(&cpu_ptr_, size_);
     }
     CUDA_CHECK(cudaMemcpy(cpu_ptr_, gpu_ptr_, size_, cudaMemcpyDeviceToHost));
     head_ = SYNCED;