blob: 11de49799358c45474ea3e2097f5a4f0f4ca57d1 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
|
#include <boost/thread.hpp>
#include <exception>
#include "caffe/internal_thread.hpp"
#include "caffe/util/math_functions.hpp"
namespace caffe {
InternalThread::~InternalThread() {
StopInternalThread();
}
bool InternalThread::is_started() const {
return thread_ && thread_->joinable();
}
bool InternalThread::must_stop() {
return thread_ && thread_->interruption_requested();
}
void InternalThread::StartInternalThread() {
CHECK(!is_started()) << "Threads should persist and not be restarted.";
int device = 0;
#ifndef CPU_ONLY
CUDA_CHECK(cudaGetDevice(&device));
#endif
Caffe::Brew mode = Caffe::mode();
int rand_seed = caffe_rng_rand();
int solver_count = Caffe::solver_count();
int solver_rank = Caffe::solver_rank();
bool multiprocess = Caffe::multiprocess();
try {
thread_.reset(new boost::thread(&InternalThread::entry, this, device, mode,
rand_seed, solver_count, solver_rank, multiprocess));
} catch (std::exception& e) {
LOG(FATAL) << "Thread exception: " << e.what();
}
}
void InternalThread::entry(int device, Caffe::Brew mode, int rand_seed,
int solver_count, int solver_rank, bool multiprocess) {
#ifndef CPU_ONLY
CUDA_CHECK(cudaSetDevice(device));
#endif
Caffe::set_mode(mode);
Caffe::set_random_seed(rand_seed);
Caffe::set_solver_count(solver_count);
Caffe::set_solver_rank(solver_rank);
Caffe::set_multiprocess(multiprocess);
InternalThreadEntry();
}
void InternalThread::StopInternalThread() {
if (is_started()) {
thread_->interrupt();
try {
thread_->join();
} catch (boost::thread_interrupted&) {
} catch (std::exception& e) {
LOG(FATAL) << "Thread exception: " << e.what();
}
}
}
} // namespace caffe
|
