Thread-local Caffe

diff --git a/include/caffe/common.hpp b/include/caffe/common.hpp
index 5f86bc2..3fa8143 100644 (file)
--- a/include/caffe/common.hpp
+++ b/include/caffe/common.hpp
@@ -98,12 +98,12 @@ void GlobalInit(int* pargc, char*** pargv);
 class Caffe {
  public:
   ~Caffe();
-  inline static Caffe& Get() {
-    if (!singleton_.get()) {
-      singleton_.reset(new Caffe());
-    }
-    return *singleton_;
-  }
+
+  // Thread local context for Caffe. Moved to common.cpp instead of
+  // including boost/thread.hpp to avoid a boost/NVCC issues (#1009, #1010)
+  // on OSX. Also fails on Linux with CUDA 7.0.18.
+  static Caffe& Get();
+
   enum Brew { CPU, GPU };
 
   // This random number generator facade hides boost and CUDA rng
@@ -158,7 +158,6 @@ class Caffe {
   shared_ptr<RNG> random_generator_;
 
   Brew mode_;
-  static shared_ptr<Caffe> singleton_;
 
  private:
   // The private constructor to avoid duplicate instantiation.

diff --git a/include/caffe/internal_thread.hpp b/include/caffe/internal_thread.hpp
index 815ca54..bcff318 100644 (file)
--- a/include/caffe/internal_thread.hpp
+++ b/include/caffe/internal_thread.hpp
@@ -14,14 +14,19 @@ namespace caffe {
 /**
  * Virtual class encapsulate boost::thread for use in base class
  * The child class will acquire the ability to run a single thread,
- * by reimplementing the virutal function InternalThreadEntry.
+ * by reimplementing the virtual function InternalThreadEntry.
  */
 class InternalThread {
  public:
   InternalThread() : thread_() {}
   virtual ~InternalThread();
 
-  /** Returns true if the thread was successfully started. **/
+  /**
+   * Caffe's thread local state will be initialized using the current
+   * thread values, e.g. device id, solver index etc. The random seed
+   * is initialized using caffe_rng_rand.
+   * Will not return until the internal thread has exited.
+   */
   bool StartInternalThread();
 
   /** Will not return until the internal thread has exited. */
@@ -34,6 +39,9 @@ class InternalThread {
       with the code you want your thread to run. */
   virtual void InternalThreadEntry() {}
 
+ private:
+  void entry(int device, Caffe::Brew mode, int rand_seed);
+
   shared_ptr<boost::thread> thread_;
 };
 

diff --git a/src/caffe/common.cpp b/src/caffe/common.cpp
index af96cac..0215c76 100644 (file)
--- a/src/caffe/common.cpp
+++ b/src/caffe/common.cpp
@@ -1,3 +1,4 @@
+#include <boost/thread.hpp>
 #include <glog/logging.h>
 #include <cstdio>
 #include <ctime>
@@ -7,7 +8,15 @@
 
 namespace caffe {
 
-shared_ptr<Caffe> Caffe::singleton_;
+// Make sure each thread can have different values.
+static boost::thread_specific_ptr<Caffe> thread_instance_;
+
+Caffe& Caffe::Get() {
+  if (!thread_instance_.get()) {
+    thread_instance_.reset(new Caffe());
+  }
+  return *(thread_instance_.get());
+}
 
 // random seeding
 int64_t cluster_seedgen(void) {

diff --git a/src/caffe/internal_thread.cpp b/src/caffe/internal_thread.cpp
index c2d19d4..2be88b3 100644 (file)
--- a/src/caffe/internal_thread.cpp
+++ b/src/caffe/internal_thread.cpp
@@ -1,9 +1,15 @@
 #include <boost/thread.hpp>
+
 #include "caffe/internal_thread.hpp"
+#include "caffe/util/math_functions.hpp"
 
 namespace caffe {
 
 InternalThread::~InternalThread() {
+  StopInternalThread();
+}
+
+InternalThread::~InternalThread() {
   WaitForInternalThreadToExit();
 }
 
@@ -11,20 +17,37 @@ bool InternalThread::is_started() const {
   return thread_.get() != NULL && thread_->joinable();
 }
 
-
 bool InternalThread::StartInternalThread() {
   if (!WaitForInternalThreadToExit()) {
     return false;
   }
+
+  int device = 0;
+#ifndef CPU_ONLY
+  CUDA_CHECK(cudaGetDevice(&device));
+#endif
+  Caffe::Brew mode = Caffe::mode();
+  int rand_seed = caffe_rng_rand();
+
   try {
-    thread_.reset(
-        new boost::thread(&InternalThread::InternalThreadEntry, this));
+    thread_.reset(new boost::thread(&InternalThread::entry, this, device, mode,
+          rand_seed));
   } catch (...) {
     return false;
   }
   return true;
 }
 
+void InternalThread::entry(int device, Caffe::Brew mode, int rand_seed) {
+#ifndef CPU_ONLY
+  CUDA_CHECK(cudaSetDevice(device));
+#endif
+  Caffe::set_mode(mode);
+  Caffe::set_random_seed(rand_seed);
+
+  InternalThreadEntry();
+}
+
 /** Will not return until the internal thread has exited. */
 bool InternalThread::WaitForInternalThreadToExit() {
   if (is_started()) {

diff --git a/src/caffe/test/test_internal_thread.cpp b/src/caffe/test/test_internal_thread.cpp
index 31882b6..390c8ed 100644 (file)
--- a/src/caffe/test/test_internal_thread.cpp
+++ b/src/caffe/test/test_internal_thread.cpp
@@ -2,6 +2,7 @@
 #include "gtest/gtest.h"
 
 #include "caffe/internal_thread.hpp"
+#include "caffe/util/math_functions.hpp"
 
 #include "caffe/test/test_caffe_main.hpp"
 
@@ -19,5 +20,34 @@ TEST_F(InternalThreadTest, TestStartAndExit) {
   EXPECT_FALSE(thread.is_started());
 }
 
+class TestThreadA : public InternalThread {
+  void InternalThreadEntry() {
+    EXPECT_EQ(4244559767, caffe_rng_rand());
+  }
+};
+
+class TestThreadB : public InternalThread {
+  void InternalThreadEntry() {
+    EXPECT_EQ(1726478280, caffe_rng_rand());
+  }
+};
+
+TEST_F(InternalThreadTest, TestRandomSeed) {
+  TestThreadA t1;
+  Caffe::set_random_seed(9658361);
+  EXPECT_TRUE(t1.StartInternalThread());
+  EXPECT_TRUE(t1.WaitForInternalThreadToExit());
+
+  TestThreadA t2;
+  Caffe::set_random_seed(9658361);
+  EXPECT_TRUE(t2.StartInternalThread());
+  EXPECT_TRUE(t2.WaitForInternalThreadToExit());
+
+  TestThreadB t3;
+  Caffe::set_random_seed(3435563);
+  EXPECT_TRUE(t3.StartInternalThread());
+  EXPECT_TRUE(t3.WaitForInternalThreadToExit());
+}
+
 }  // namespace caffe