Skip to content

Commit 6d09ca2

Browse files
committed
Speeding up the GPU solvers
1 parent b086cc3 commit 6d09ca2

12 files changed

Lines changed: 223 additions & 161 deletions

src/caffe/solvers/adadelta_solver.cpp

Lines changed: 11 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,12 @@ void AdaDeltaSolver<Dtype>::AdaDeltaPreSolve() {
1616
}
1717
}
1818

19+
#ifndef CPU_ONLY
20+
template <typename Dtype>
21+
void adadelta_update_gpu(int N, Dtype* g, Dtype* h, Dtype* h2, Dtype momentum,
22+
Dtype delta, Dtype local_rate);
23+
#endif
24+
1925
template <typename Dtype>
2026
void AdaDeltaSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
2127
const vector<Blob<Dtype>*>& net_params = this->net_->learnable_params();
@@ -85,61 +91,11 @@ void AdaDeltaSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
8591
}
8692
case Caffe::GPU: {
8793
#ifndef CPU_ONLY
88-
// compute square of gradient in update
89-
caffe_gpu_powx(net_params[param_id]->count(),
90-
net_params[param_id]->gpu_diff(), Dtype(2),
91-
this->update_[param_id]->mutable_gpu_data());
92-
93-
// update history of gradients
94-
caffe_gpu_axpby(net_params[param_id]->count(), Dtype(1) - momentum,
95-
this->update_[param_id]->gpu_data(), momentum,
96-
this->history_[param_id]->mutable_gpu_data());
97-
98-
// add delta to history to guard against dividing by zero later
99-
caffe_gpu_set(net_params[param_id]->count(), delta,
100-
this->temp_[param_id]->mutable_gpu_data());
101-
102-
caffe_gpu_add(net_params[param_id]->count(),
103-
this->temp_[param_id]->gpu_data(),
104-
this->history_[update_history_offset + param_id]->gpu_data(),
105-
this->update_[param_id]->mutable_gpu_data());
106-
107-
caffe_gpu_add(net_params[param_id]->count(),
108-
this->temp_[param_id]->gpu_data(),
109-
this->history_[param_id]->gpu_data(),
110-
this->temp_[param_id]->mutable_gpu_data());
111-
112-
// divide history of updates by history of gradients
113-
caffe_gpu_div(net_params[param_id]->count(),
114-
this->update_[param_id]->gpu_data(),
115-
this->temp_[param_id]->gpu_data(),
116-
this->update_[param_id]->mutable_gpu_data());
117-
118-
// jointly compute the RMS of both for update and gradient history
119-
caffe_gpu_powx(net_params[param_id]->count(),
120-
this->update_[param_id]->gpu_data(), Dtype(0.5),
121-
this->update_[param_id]->mutable_gpu_data());
122-
123-
// compute the update and copy to net_diff
124-
caffe_gpu_mul(net_params[param_id]->count(),
125-
net_params[param_id]->gpu_diff(),
126-
this->update_[param_id]->gpu_data(),
127-
net_params[param_id]->mutable_gpu_diff());
128-
129-
// compute square of update
130-
caffe_gpu_powx(net_params[param_id]->count(),
131-
net_params[param_id]->gpu_diff(), Dtype(2),
132-
this->update_[param_id]->mutable_gpu_data());
133-
134-
// update history of updates
135-
caffe_gpu_axpby(net_params[param_id]->count(), Dtype(1) - momentum,
136-
this->update_[param_id]->gpu_data(), momentum,
137-
this->history_[update_history_offset + param_id]->mutable_gpu_data());
138-
139-
// apply learning rate
140-
caffe_gpu_scale(net_params[param_id]->count(), local_rate,
141-
net_params[param_id]->gpu_diff(),
142-
net_params[param_id]->mutable_gpu_diff());
94+
adadelta_update_gpu(net_params[param_id]->count(),
95+
net_params[param_id]->mutable_gpu_diff(),
96+
this->history_[param_id]->mutable_gpu_data(),
97+
this->history_[update_history_offset + param_id]->mutable_gpu_data(),
98+
momentum, delta, local_rate);
14399
#else
144100
NO_GPU;
145101
#endif
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
#include "caffe/util/math_functions.hpp"
2+
3+
4+
namespace caffe {
5+
6+
template <typename Dtype>
7+
__global__ void AdaDeltaUpdate(int N, Dtype* g, Dtype* h, Dtype* h2,
8+
Dtype momentum, Dtype delta, Dtype local_rate) {
9+
CUDA_KERNEL_LOOP(i, N) {
10+
float gi = g[i];
11+
float hi = h[i] = momentum * h[i] + (1-momentum) * gi * gi;
12+
gi = gi * sqrt((h2[i] + delta) / (hi + delta));
13+
h2[i] = momentum * h2[i] + (1-momentum) * gi * gi;
14+
g[i] = local_rate * gi;
15+
}
16+
}
17+
template <typename Dtype>
18+
void adadelta_update_gpu(int N, Dtype* g, Dtype* h, Dtype* h2, Dtype momentum,
19+
Dtype delta, Dtype local_rate) {
20+
AdaDeltaUpdate<Dtype> // NOLINT_NEXT_LINE(whitespace/operators)
21+
<<<CAFFE_GET_BLOCKS(N), CAFFE_CUDA_NUM_THREADS>>>(
22+
N, g, h, h2, momentum, delta, local_rate);
23+
CUDA_POST_KERNEL_CHECK;
24+
}
25+
template void adadelta_update_gpu<float>(int , float*, float*, float*,
26+
float, float, float);
27+
template void adadelta_update_gpu<double>(int, double*, double*, double*,
28+
double, double, double);
29+
30+
} // namespace caffe

src/caffe/solvers/adagrad_solver.cpp

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,12 @@
44

55
namespace caffe {
66

7+
#ifndef CPU_ONLY
8+
template <typename Dtype>
9+
void adagrad_update_gpu(int N, Dtype* g, Dtype* h, Dtype delta,
10+
Dtype local_rate);
11+
#endif
12+
713
template <typename Dtype>
814
void AdaGradSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
915
CHECK(Caffe::root_solver());
@@ -45,34 +51,9 @@ void AdaGradSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
4551
}
4652
case Caffe::GPU: {
4753
#ifndef CPU_ONLY
48-
// compute square of gradient in update
49-
caffe_gpu_powx(net_params[param_id]->count(),
50-
net_params[param_id]->gpu_diff(), Dtype(2),
51-
this->update_[param_id]->mutable_gpu_data());
52-
53-
// update history
54-
caffe_gpu_add(net_params[param_id]->count(),
55-
this->update_[param_id]->gpu_data(),
56-
this->history_[param_id]->gpu_data(),
57-
this->history_[param_id]->mutable_gpu_data());
58-
59-
// prepare update
60-
caffe_gpu_powx(net_params[param_id]->count(),
61-
this->history_[param_id]->gpu_data(), Dtype(0.5),
62-
this->update_[param_id]->mutable_gpu_data());
63-
64-
caffe_gpu_add_scalar(net_params[param_id]->count(),
65-
delta, this->update_[param_id]->mutable_gpu_data());
66-
67-
caffe_gpu_div(net_params[param_id]->count(),
68-
net_params[param_id]->gpu_diff(),
69-
this->update_[param_id]->gpu_data(),
70-
this->update_[param_id]->mutable_gpu_data());
71-
72-
// scale and copy
73-
caffe_gpu_axpby(net_params[param_id]->count(), local_rate,
74-
this->update_[param_id]->gpu_data(), Dtype(0),
75-
net_params[param_id]->mutable_gpu_diff());
54+
adagrad_update_gpu(net_params[param_id]->count(),
55+
net_params[param_id]->mutable_gpu_diff(),
56+
this->history_[param_id]->mutable_gpu_data(), delta, local_rate);
7657
#else
7758
NO_GPU;
7859
#endif
Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#include "caffe/util/math_functions.hpp"
2+
3+
4+
namespace caffe {
5+
6+
template <typename Dtype>
7+
__global__ void AdaGradUpdate(int N, Dtype* g, Dtype* h, Dtype delta,
8+
Dtype local_rate) {
9+
CUDA_KERNEL_LOOP(i, N) {
10+
float gi = g[i];
11+
float hi = h[i] = h[i] + gi*gi;
12+
g[i] = local_rate * gi / (sqrt(hi) + delta);
13+
}
14+
}
15+
template <typename Dtype>
16+
void adagrad_update_gpu(int N, Dtype* g, Dtype* h, Dtype delta,
17+
Dtype local_rate) {
18+
AdaGradUpdate<Dtype> // NOLINT_NEXT_LINE(whitespace/operators)
19+
<<<CAFFE_GET_BLOCKS(N), CAFFE_CUDA_NUM_THREADS>>>(
20+
N, g, h, delta, local_rate);
21+
CUDA_POST_KERNEL_CHECK;
22+
}
23+
template void adagrad_update_gpu<float>(int, float*, float*, float, float);
24+
template void adagrad_update_gpu<double>(int, double*, double*, double, double);
25+
26+
} // namespace caffe

src/caffe/solvers/adam_solver.cpp

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,12 @@ void AdamSolver<Dtype>::AdamPreSolve() {
1616
}
1717
}
1818

19+
#ifndef CPU_ONLY
20+
template <typename Dtype>
21+
void adam_update_gpu(int N, Dtype* g, Dtype* m, Dtype* v, Dtype beta1,
22+
Dtype beta2, Dtype eps_hat, Dtype corrected_local_rate);
23+
#endif
24+
1925
template <typename Dtype>
2026
void AdamSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
2127
const vector<Blob<Dtype>*>& net_params = this->net_->learnable_params();
@@ -69,34 +75,9 @@ void AdamSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
6975
}
7076
case Caffe::GPU: {
7177
#ifndef CPU_ONLY
72-
// update m <- \beta_1 m_{t-1} + (1-\beta_1)g_t
73-
caffe_gpu_axpby(N, Dtype(1)-beta1,
74-
net_params[param_id]->gpu_diff(), beta1,
75-
val_m->mutable_gpu_data());
76-
77-
// update v <- \beta_2 m_{t-1} + (1-\beta_2)g_t^2
78-
caffe_gpu_mul(N,
79-
net_params[param_id]->gpu_diff(),
80-
net_params[param_id]->gpu_diff(),
81-
val_t->mutable_gpu_data());
82-
caffe_gpu_axpby(N, Dtype(1)-beta2,
83-
val_t->gpu_data(), beta2,
84-
val_v->mutable_gpu_data());
85-
86-
// set update
87-
caffe_gpu_powx(N,
88-
val_v->gpu_data(), Dtype(0.5),
89-
val_t->mutable_gpu_data());
90-
caffe_gpu_add_scalar(N, eps_hat,
91-
val_t->mutable_gpu_data());
92-
caffe_gpu_div(N,
93-
val_m->gpu_data(),
94-
val_t->gpu_data(),
95-
val_t->mutable_gpu_data());
96-
97-
caffe_gpu_scale(N, local_rate*correction,
98-
val_t->gpu_data(),
99-
net_params[param_id]->mutable_gpu_diff());
78+
adam_update_gpu(N, net_params[param_id]->mutable_gpu_diff(),
79+
val_m->mutable_gpu_data(), val_v->mutable_gpu_data(), beta1, beta2,
80+
eps_hat, local_rate*correction);
10081
#else
10182
NO_GPU;
10283
#endif

src/caffe/solvers/adam_solver.cu

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
#include "caffe/util/math_functions.hpp"
2+
3+
4+
namespace caffe {
5+
6+
template <typename Dtype>
7+
__global__ void AdamUpdate(int N, Dtype* g, Dtype* m, Dtype* v,
8+
Dtype beta1, Dtype beta2, Dtype eps_hat, Dtype corrected_local_rate) {
9+
CUDA_KERNEL_LOOP(i, N) {
10+
float gi = g[i];
11+
float mi = m[i] = m[i]*beta1 + gi*(1-beta1);
12+
float vi = v[i] = v[i]*beta2 + gi*gi*(1-beta2);
13+
g[i] = corrected_local_rate * mi / (sqrt(vi) + eps_hat);
14+
}
15+
}
16+
template <typename Dtype>
17+
void adam_update_gpu(int N, Dtype* g, Dtype* m, Dtype* v, Dtype beta1,
18+
Dtype beta2, Dtype eps_hat, Dtype corrected_local_rate) {
19+
AdamUpdate<Dtype> // NOLINT_NEXT_LINE(whitespace/operators)
20+
<<<CAFFE_GET_BLOCKS(N), CAFFE_CUDA_NUM_THREADS>>>(
21+
N, g, m, v, beta1, beta2, eps_hat, corrected_local_rate);
22+
CUDA_POST_KERNEL_CHECK;
23+
}
24+
template void adam_update_gpu<float>(int, float*, float*, float*,
25+
float, float, float, float);
26+
template void adam_update_gpu<double>(int, double*, double*, double*,
27+
double, double, double, double);
28+
29+
} // namespace caffe

src/caffe/solvers/nesterov_solver.cpp

Lines changed: 10 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,12 @@
44

55
namespace caffe {
66

7+
#ifndef CPU_ONLY
8+
template <typename Dtype>
9+
void nesterov_update_gpu(int N, Dtype* g, Dtype* h, Dtype momentum,
10+
Dtype local_rate);
11+
#endif
12+
713
template <typename Dtype>
814
void NesterovSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
915
CHECK(Caffe::root_solver());
@@ -36,25 +42,10 @@ void NesterovSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
3642
}
3743
case Caffe::GPU: {
3844
#ifndef CPU_ONLY
39-
// save history momentum for stepping back
40-
caffe_copy(net_params[param_id]->count(),
41-
this->history_[param_id]->gpu_data(),
42-
this->update_[param_id]->mutable_gpu_data());
43-
44-
// update history
45-
caffe_gpu_axpby(net_params[param_id]->count(), local_rate,
46-
net_params[param_id]->gpu_diff(), momentum,
47-
this->history_[param_id]->mutable_gpu_data());
48-
49-
// compute update: step back then over step
50-
caffe_gpu_axpby(net_params[param_id]->count(), Dtype(1) + momentum,
51-
this->history_[param_id]->gpu_data(), -momentum,
52-
this->update_[param_id]->mutable_gpu_data());
53-
54-
// copy
55-
caffe_copy(net_params[param_id]->count(),
56-
this->update_[param_id]->gpu_data(),
57-
net_params[param_id]->mutable_gpu_diff());
45+
nesterov_update_gpu(net_params[param_id]->count(),
46+
net_params[param_id]->mutable_gpu_diff(),
47+
this->history_[param_id]->mutable_gpu_data(),
48+
momentum, local_rate);
5849
#else
5950
NO_GPU;
6051
#endif
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
#include "caffe/util/math_functions.hpp"
2+
3+
4+
namespace caffe {
5+
6+
template <typename Dtype>
7+
__global__ void NesterovUpdate(int N, Dtype* g, Dtype* h,
8+
Dtype momentum, Dtype local_rate) {
9+
CUDA_KERNEL_LOOP(i, N) {
10+
float hi = h[i];
11+
float hi_new = h[i] = momentum * hi + local_rate * g[i];
12+
g[i] = (1+momentum) * hi_new - momentum * hi;
13+
}
14+
}
15+
template <typename Dtype>
16+
void nesterov_update_gpu(int N, Dtype* g, Dtype* h, Dtype momentum,
17+
Dtype local_rate) {
18+
NesterovUpdate<Dtype> // NOLINT_NEXT_LINE(whitespace/operators)
19+
<<<CAFFE_GET_BLOCKS(N), CAFFE_CUDA_NUM_THREADS>>>(
20+
N, g, h, momentum, local_rate);
21+
CUDA_POST_KERNEL_CHECK;
22+
}
23+
template void nesterov_update_gpu<float>(int, float*, float*, float, float);
24+
template void nesterov_update_gpu<double>(int, double*, double*, double,
25+
double);
26+
27+
} // namespace caffe

src/caffe/solvers/rmsprop_solver.cpp

Lines changed: 10 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,12 @@
44

55
namespace caffe {
66

7+
#ifndef CPU_ONLY
8+
template <typename Dtype>
9+
void rmsprop_update_gpu(int N, Dtype* g, Dtype* h, Dtype rms_decay,
10+
Dtype delta, Dtype local_rate);
11+
#endif
12+
713
template <typename Dtype>
814
void RMSPropSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
915
const vector<Blob<Dtype>*>& net_params = this->net_->learnable_params();
@@ -45,31 +51,10 @@ void RMSPropSolver<Dtype>::ComputeUpdateValue(int param_id, Dtype rate) {
4551
break;
4652
case Caffe::GPU:
4753
#ifndef CPU_ONLY
48-
// compute square of gradient in update
49-
caffe_gpu_powx(net_params[param_id]->count(),
50-
net_params[param_id]->gpu_diff(), Dtype(2),
51-
this->update_[param_id]->mutable_gpu_data());
52-
53-
// update history
54-
caffe_gpu_axpby(net_params[param_id] -> count(),
55-
Dtype(1-rms_decay), this->update_[param_id]->gpu_data(),
56-
rms_decay, this->history_[param_id]-> mutable_gpu_data());
57-
58-
// prepare update
59-
caffe_gpu_powx(net_params[param_id]->count(),
60-
this->history_[param_id]->gpu_data(), Dtype(0.5),
61-
this->update_[param_id]->mutable_gpu_data());
62-
63-
caffe_gpu_add_scalar(net_params[param_id]->count(),
64-
delta, this->update_[param_id]->mutable_gpu_data());
65-
66-
caffe_gpu_div(net_params[param_id]->count(),
67-
net_params[param_id]->gpu_diff(), this->update_[param_id]->gpu_data(),
68-
this->update_[param_id]->mutable_gpu_data());
69-
70-
caffe_gpu_axpby(net_params[param_id]->count(), local_rate,
71-
this->update_[param_id]->gpu_data(), Dtype(0),
72-
net_params[param_id]->mutable_gpu_diff());
54+
rmsprop_update_gpu(net_params[param_id]->count(),
55+
net_params[param_id]->mutable_gpu_diff(),
56+
this->history_[param_id]->mutable_gpu_data(),
57+
rms_decay, delta, local_rate);
7358
#else
7459
NO_GPU;
7560
#endif

0 commit comments

Comments
 (0)