internal_thread.hpp 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. #ifndef CAFFE_INTERNAL_THREAD_HPP_
  2. #define CAFFE_INTERNAL_THREAD_HPP_
  3. #include "caffe/common.hpp"
  4. /**
  5. Forward declare boost::thread instead of including boost/thread.hpp
  6. to avoid a boost/NVCC issues (#1009, #1010) on OSX.
  7. */
  8. namespace boost { class thread; }
  9. namespace caffe {
  10. /**
  11. * Virtual class encapsulate boost::thread for use in base class
  12. * The child class will acquire the ability to run a single thread,
  13. * by reimplementing the virtual function InternalThreadEntry.
  14. */
  15. class InternalThread {
  16. public:
  17. InternalThread() : thread_() {}
  18. virtual ~InternalThread();
  19. /**
  20. * Caffe's thread local state will be initialized using the current
  21. * thread values, e.g. device id, solver index etc. The random seed
  22. * is initialized using caffe_rng_rand.
  23. */
  24. void StartInternalThread();
  25. /** Will not return until the internal thread has exited. */
  26. void StopInternalThread();
  27. bool is_started() const;
  28. protected:
  29. /* Implement this method in your subclass
  30. with the code you want your thread to run. */
  31. virtual void InternalThreadEntry() {}
  32. /* Should be tested when running loops to exit when requested. */
  33. bool must_stop();
  34. private:
  35. void entry(int device, Caffe::Brew mode, int rand_seed,
  36. int solver_count, int solver_rank, bool multiprocess);
  37. shared_ptr<boost::thread> thread_;
  38. };
  39. } // namespace caffe
  40. #endif // CAFFE_INTERNAL_THREAD_HPP_