| 311 | } |
| 312 | |
| 313 | Tensor CpuConvBackwardW(const Tensor &dy, const Tensor &x, const Tensor &W, |
| 314 | const ConvHandle &ch) { |
| 315 | CHECK_EQ(dy.device()->lang(), kCpp); |
| 316 | CHECK_EQ(x.device()->lang(), kCpp); |
| 317 | CHECK_EQ(W.device()->lang(), kCpp); |
| 318 | |
| 319 | CHECK(dy.shape(1) == ch.num_filters && dy.shape(2) == ch.conv_height && |
| 320 | dy.shape(3) == ch.conv_width) |
| 321 | << "input gradients shape should not change"; |
| 322 | |
| 323 | CHECK(x.shape(1) == ch.channels && x.shape(2) == ch.height && |
| 324 | x.shape(3) == ch.width) |
| 325 | << "input sample shape should not change"; |
| 326 | |
| 327 | #ifdef USE_DNNL |
| 328 | Tensor dW; |
| 329 | dW.ResetLike(W); |
| 330 | |
| 331 | dy.device()->Exec( |
| 332 | [dy, dW, x, &W, &ch](Context *ctx) mutable { |
| 333 | using namespace dnnl; |
| 334 | auto eng = ctx->dnnl_engine; |
| 335 | auto s = ctx->dnnl_stream; |
| 336 | using tag = memory::format_tag; |
| 337 | auto dtype = dnnl::memory::data_type::f32; |
| 338 | |
| 339 | auto conv_src_md = memory::desc({ch.x_dims}, dtype, tag::nchw); |
| 340 | auto conv_weights_md = memory::desc({ch.w_dims}, dtype, tag::goihw); |
| 341 | auto conv_bias_md = memory::desc({ch.b_dims}, dtype, tag::x); |
| 342 | auto conv_dst_md = memory::desc({ch.o_dims}, dtype, tag::nchw); |
| 343 | |
| 344 | auto conv_user_src_memory = |
| 345 | memory(conv_src_md, eng, x.block()->mutable_data()); |
| 346 | auto conv_user_diff_weights_memory = |
| 347 | memory(conv_weights_md, eng, dW.block()->mutable_data()); |
| 348 | auto conv_diff_bias_memory = |
| 349 | memory(conv_bias_md, eng, ch.db->block()->mutable_data()); |
| 350 | auto conv_user_diff_dst_memory = |
| 351 | memory(conv_dst_md, eng, dy.block()->mutable_data()); |
| 352 | |
| 353 | auto conv_desc = convolution_forward::desc( |
| 354 | prop_kind::forward, algorithm::convolution_direct, conv_src_md, |
| 355 | conv_weights_md, conv_bias_md, conv_dst_md, ch.s_dims, ch.p_dims, |
| 356 | ch.p_dims); |
| 357 | auto conv_pd = convolution_forward::primitive_desc(conv_desc, eng); |
| 358 | |
| 359 | // auto conv_pd = *ch.conv_pd; // very slow |
| 360 | |
| 361 | auto conv_bwd_src_memory = conv_user_src_memory; |
| 362 | auto conv_diff_weights_memory = conv_user_diff_weights_memory; |
| 363 | auto conv_diff_dst_memory = conv_user_diff_dst_memory; |
| 364 | |
| 365 | auto conv_bwd_weights_desc = convolution_backward_weights::desc( |
| 366 | algorithm::convolution_direct, conv_src_md, conv_weights_md, |
| 367 | conv_bias_md, conv_dst_md, ch.s_dims, ch.p_dims, ch.p_dims); |
| 368 | auto conv_bwd_weights_pd = convolution_backward_weights::primitive_desc( |
| 369 | conv_bwd_weights_desc, eng, conv_pd); |
| 370 | |