| 583 | } |
| 584 | |
| 585 | Tensor GpuConvForward(const Tensor &x, const Tensor &W, const Tensor &b, |
| 586 | const CudnnConvHandle &cch) { |
| 587 | CHECK_EQ(x.device()->lang(), kCuda); |
| 588 | CHECK(x.shape(1) == cch.channels && x.shape(2) == cch.height && |
| 589 | x.shape(3) == cch.width) |
| 590 | << "input sample shape should not change"; |
| 591 | |
| 592 | CHECK(W.shape(0) == cch.num_filters && |
| 593 | W.shape(1) == cch.channels_per_filter && W.shape(2) == cch.kernel_h && |
| 594 | W.shape(3) == cch.kernel_w) |
| 595 | << "weights shape should not change"; |
| 596 | |
| 597 | DataType dtype = x.data_type(); |
| 598 | auto dev = x.device(); |
| 599 | |
| 600 | Shape shape{cch.batchsize, cch.num_filters, cch.conv_height, cch.conv_width}; |
| 601 | Tensor output(shape, dev, dtype); |
| 602 | |
| 603 | output.device()->Exec( |
| 604 | [output, x, &W, &cch](Context *ctx) mutable { |
| 605 | Block *inblock = x.block(), *outblock = output.block(), |
| 606 | *wblock = W.block(); |
| 607 | float alpha = 1.f, beta = 0.f; |
| 608 | cudnnConvolutionForward(ctx->cudnn_handle, &alpha, cch.x_desc, |
| 609 | inblock->data(), cch.filter_desc, |
| 610 | wblock->data(), cch.conv_desc, cch.fp_alg, |
| 611 | cch.workspace.block()->mutable_data(), |
| 612 | cch.workspace_count * SizeOf(x.data_type()), &beta, |
| 613 | cch.y_desc, outblock->mutable_data()); |
| 614 | }, |
| 615 | {x.block(), W.block()}, {output.block(), cch.workspace.block()}, |
| 616 | "cudnnConvForward"); |
| 617 | |
| 618 | if (cch.bias_term) { |
| 619 | Tensor outputFake(output); |
| 620 | output.device()->Exec( |
| 621 | [output, outputFake, &b, &cch](Context *ctx) mutable { |
| 622 | float beta = 1.f, alpha = 1.0f; |
| 623 | Block *outblock = output.block(), *bblock = b.block(); |
| 624 | cudnnAddTensor(ctx->cudnn_handle, &alpha, cch.bias_desc, |
| 625 | bblock->data(), &beta, cch.y_desc, |
| 626 | outblock->mutable_data()); |
| 627 | }, |
| 628 | {output.block(), b.block()}, {output.block()}, "cudnnAddTensor"); |
| 629 | } |
| 630 | |
| 631 | return output; |
| 632 | } |
| 633 | |
| 634 | Tensor GpuConvBackwardx(const Tensor &dy, const Tensor &W, const Tensor &x, |
| 635 | const CudnnConvHandle &cch) { |