| 323 | |
| 324 | template<typename Dtype> |
| 325 | void P2PSync<Dtype>::on_gradients_ready() { |
| 326 | #ifndef CPU_ONLY |
| 327 | #ifdef DEBUG |
| 328 | int device; |
| 329 | CUDA_CHECK(cudaGetDevice(&device)); |
| 330 | CHECK(device == solver_->param().device_id()); |
| 331 | #endif |
| 332 | |
| 333 | // Sum children gradients as they appear in the queue |
| 334 | for (int i = 0; i < children_.size(); ++i) { |
| 335 | P2PSync<Dtype> *child = queue_.pop(); |
| 336 | Dtype* src = child->parent_grads_; |
| 337 | Dtype* dst = diff_; |
| 338 | |
| 339 | #ifdef DEBUG |
| 340 | bool ok = false; |
| 341 | for (int j = 0; j < children_.size(); ++j) { |
| 342 | if (child == children_[j]) { |
| 343 | ok = true; |
| 344 | } |
| 345 | } |
| 346 | CHECK(ok); |
| 347 | cudaPointerAttributes attributes; |
| 348 | CUDA_CHECK(cudaPointerGetAttributes(&attributes, src)); |
| 349 | CHECK(attributes.device == device); |
| 350 | CUDA_CHECK(cudaPointerGetAttributes(&attributes, dst)); |
| 351 | CHECK(attributes.device == device); |
| 352 | #endif |
| 353 | |
| 354 | caffe_gpu_add(size_, src, dst, dst); |
| 355 | } |
| 356 | |
| 357 | // Send gradients to parent |
| 358 | if (parent_) { |
| 359 | Dtype* src = diff_; |
| 360 | Dtype* dst = parent_grads_; |
| 361 | |
| 362 | #ifdef DEBUG |
| 363 | cudaPointerAttributes attributes; |
| 364 | CUDA_CHECK(cudaPointerGetAttributes(&attributes, src)); |
| 365 | CHECK(attributes.device == device); |
| 366 | CUDA_CHECK(cudaPointerGetAttributes(&attributes, dst)); |
| 367 | CHECK(attributes.device == parent_->solver_->param().device_id()); |
| 368 | #endif |
| 369 | |
| 370 | CUDA_CHECK(cudaMemcpyAsync(dst, src, size_ * sizeof(Dtype), // |
| 371 | cudaMemcpyDeviceToDevice, cudaStreamDefault)); |
| 372 | CUDA_CHECK(cudaStreamSynchronize(cudaStreamDefault)); |
| 373 | parent_->queue_.push(this); |
| 374 | } else { |
| 375 | // Loss functions divide gradients by the batch size, so to compensate |
| 376 | // for split batch, the root solver divides by number of solvers. |
| 377 | caffe_gpu_scal(size_, Dtype(1.0 / Caffe::solver_count()), diff_); |
| 378 | } |
| 379 | #endif |
| 380 | } |
| 381 | |
| 382 | template<typename Dtype> |