| 212 | } |
| 213 | |
| 214 | const std::pair<Tensor, vector<Tensor>> CudnnConvolution::Backward( |
| 215 | int flag, const Tensor &grad) { |
| 216 | CHECK(has_init_cudnn_); |
| 217 | CHECK_EQ(grad.device()->lang(), kCuda); |
| 218 | CHECK_EQ(grad.nDim(), 4u); |
| 219 | CHECK(!buf_.empty()); |
| 220 | Tensor src_data = buf_.top(); |
| 221 | buf_.pop(); |
| 222 | vector<Tensor> param_grad; |
| 223 | Tensor dx; |
| 224 | dx.ResetLike(src_data); |
| 225 | Tensor db, dw; |
| 226 | dw.ResetLike(weight_); |
| 227 | |
| 228 | // LOG(ERROR) << "backward bias"; |
| 229 | if (bias_term_) { |
| 230 | db.ResetLike(bias_); |
| 231 | dx.device()->Exec([grad, db, this](Context * ctx) { |
| 232 | Block *dyblock = grad.block(), *dbblock = db.block(); |
| 233 | float alpha = 1.f, beta = 0.f; |
| 234 | cudnnConvolutionBackwardBias(ctx->cudnn_handle, &alpha, this->y_desc_, |
| 235 | dyblock->data(), &beta, this->bias_desc_, |
| 236 | dbblock->mutable_data()); |
| 237 | }, {grad.block()}, {db.block()}); |
| 238 | } |
| 239 | // LOG(ERROR) << "backward w"; |
| 240 | dx.device()->Exec([grad, dw, src_data, this](Context * ctx) { |
| 241 | Block *inblock = src_data.block(), *dyblock = grad.block(), |
| 242 | *dwblock = dw.block(); |
| 243 | float alpha = 1.f, beta = 0.f; |
| 244 | cudnnConvolutionBackwardFilter( |
| 245 | ctx->cudnn_handle, &alpha, this->x_desc_, inblock->data(), |
| 246 | this->y_desc_, dyblock->data(), this->conv_desc_, this->bp_filter_alg_, |
| 247 | this->workspace_.block()->mutable_data(), |
| 248 | this->workspace_count_ * sizeof(float), &beta, this->filter_desc_, |
| 249 | dwblock->mutable_data()); |
| 250 | }, {grad.block(), src_data.block()}, {dw.block(), workspace_.block()}); |
| 251 | |
| 252 | // LOG(ERROR) << "backward src"; |
| 253 | dx.device()->Exec([dx, grad, this](Context * ctx) { |
| 254 | Block *wblock = this->weight_.block(), *dyblock = grad.block(), |
| 255 | *dxblock = dx.block(); |
| 256 | float alpha = 1.f, beta = 0.f; |
| 257 | cudnnConvolutionBackwardData(ctx->cudnn_handle, &alpha, this->filter_desc_, |
| 258 | wblock->data(), this->y_desc_, dyblock->data(), |
| 259 | this->conv_desc_, this->bp_data_alg_, |
| 260 | this->workspace_.block()->mutable_data(), |
| 261 | this->workspace_count_ * sizeof(float), &beta, |
| 262 | this->x_desc_, dxblock->mutable_data()); |
| 263 | }, {grad.block(), weight_.block()}, {dx.block(), workspace_.block()}); |
| 264 | param_grad.push_back(dw); |
| 265 | if (bias_term_) |
| 266 | param_grad.push_back(db); |
| 267 | return std::make_pair(dx, param_grad); |
| 268 | } |
| 269 | |
| 270 | } // namespace singa |
| 271 | #endif // USE_CUDNN |