| 185 | } |
| 186 | |
| 187 | Tensor GpuPoolingForward(const CudnnPoolingHandle &cph, const Tensor &x) { |
| 188 | CHECK_EQ(x.device()->lang(), kCuda); |
| 189 | CHECK_EQ(x.nDim(), 4u); |
| 190 | |
| 191 | Tensor output = Tensor( |
| 192 | Shape({cph.batchsize, cph.channels, cph.pooled_height, cph.pooled_width}), |
| 193 | x.device(), x.data_type()); |
| 194 | |
| 195 | output.device()->Exec( |
| 196 | [output, x, &cph](Context *ctx) mutable { |
| 197 | float alpha = 1.0f, beta = 0.0f; |
| 198 | cudnnPoolingForward(ctx->cudnn_handle, cph.pool_desc, &alpha, |
| 199 | cph.x_desc, x.block()->data(), &beta, cph.y_desc, |
| 200 | output.block()->mutable_data()); |
| 201 | }, |
| 202 | {x.block()}, {output.block()}, "GpuPoolingForward"); |
| 203 | |
| 204 | return output; |
| 205 | } |
| 206 | |
| 207 | Tensor GpuPoolingBackward(const CudnnPoolingHandle &cph, const Tensor &dy, |
| 208 | const Tensor &x, const Tensor &y) { |