Skip to content

Commit

Permalink
Thread-local Caffe
Browse files Browse the repository at this point in the history
  • Loading branch information
cypof committed May 19, 2015
1 parent 17e57ae commit 95dcc20
Show file tree
Hide file tree
Showing 5 changed files with 87 additions and 14 deletions.
13 changes: 6 additions & 7 deletions include/caffe/common.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down
17 changes: 14 additions & 3 deletions include/caffe/internal_thread.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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_() {}
InternalThread();
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. */
Expand All @@ -34,7 +39,13 @@ class InternalThread {
with the code you want your thread to run. */
virtual void InternalThreadEntry() {}

private:
void entry();

shared_ptr<boost::thread> thread_;
int device_;
Caffe::Brew mode_;
int rand_seed_;
};

} // namespace caffe
Expand Down
11 changes: 10 additions & 1 deletion src/caffe/common.cpp
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#include <boost/thread.hpp>
#include <glog/logging.h>
#include <cstdio>
#include <ctime>
Expand All @@ -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) {
Expand Down
30 changes: 27 additions & 3 deletions src/caffe/internal_thread.cpp
Original file line number Diff line number Diff line change
@@ -1,8 +1,17 @@
#include <boost/thread.hpp>

#include "caffe/internal_thread.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

InternalThread::InternalThread()
: thread_(),
device_(),
mode_(),
rand_seed_() {
}

InternalThread::~InternalThread() {
WaitForInternalThreadToExit();
}
Expand All @@ -11,20 +20,35 @@ bool InternalThread::is_started() const {
return thread_.get() != NULL && thread_->joinable();
}


bool InternalThread::StartInternalThread() {
if (!WaitForInternalThreadToExit()) {
return false;
}

#ifndef CPU_ONLY
CUDA_CHECK(cudaGetDevice(&device_));
#endif
mode_ = Caffe::mode();
rand_seed_ = caffe_rng_rand();

try {
thread_.reset(
new boost::thread(&InternalThread::InternalThreadEntry, this));
thread_.reset(new boost::thread(&InternalThread::entry, this));
} catch (...) {
return false;
}
return true;
}

void InternalThread::entry() {
#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()) {
Expand Down
30 changes: 30 additions & 0 deletions src/caffe/test/test_internal_thread.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand All @@ -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

0 comments on commit 95dcc20

Please sign in to comment.