| 159 | } |
| 160 | |
| 161 | const Tensor CudnnConvolution::Forward(int flag, const Tensor &input) { |
| 162 | CHECK(buf_.empty()); |
| 163 | CHECK_EQ(input.device()->lang(), kCuda); |
| 164 | CHECK_EQ(input.nDim(), 4u); |
| 165 | if (flag & kTrain) buf_.push(input); // buffer the input for backward |
| 166 | size_t batchsize = input.shape()[0]; |
| 167 | DataType dtype = input.data_type(); |
| 168 | auto dev = input.device(); |
| 169 | |
| 170 | if (!has_init_cudnn_) { |
| 171 | InitCudnn(input); |
| 172 | } else { |
| 173 | int n, c, h, w, s; |
| 174 | cudnnDataType_t type; |
| 175 | CUDNN_CHECK(cudnnGetTensor4dDescriptor(x_desc_, &type, &n, &c, &h, &w, |
| 176 | &s, &s, &s, &s)); |
| 177 | if (batchsize != static_cast<size_t>(n)) |
| 178 | InitCudnn(input); |
| 179 | CHECK(input.shape(1) == static_cast<size_t>(c) |
| 180 | && input.shape(2) == static_cast<size_t>(h) |
| 181 | && input.shape(3) == static_cast<size_t>(w)) |
| 182 | << "input sample shape should not change" |
| 183 | << "previous shape " << c << ", " << h << ", " << w |
| 184 | << "current shape " << input.shape(1) << ", " << input.shape(2) << ", " |
| 185 | << input.shape(3); |
| 186 | } |
| 187 | |
| 188 | Shape shape{batchsize, num_filters_, conv_height_, conv_width_}; |
| 189 | Tensor output(shape, dev, dtype); |
| 190 | output.device()->Exec([input, output, this](Context * ctx) { |
| 191 | Block *inblock = input.block(), *outblock = output.block(), |
| 192 | *wblock = this->weight_.block(); |
| 193 | float alpha = 1.f, beta = 0.f; |
| 194 | cudnnConvolutionForward(ctx->cudnn_handle, &alpha, this->x_desc_, |
| 195 | inblock->data(), this->filter_desc_, wblock->data(), |
| 196 | this->conv_desc_, this->fp_alg_, |
| 197 | this->workspace_.block()->mutable_data(), |
| 198 | this->workspace_count_ * sizeof(float), &beta, |
| 199 | this->y_desc_, outblock->mutable_data()); |
| 200 | }, {input.block(), weight_.block()}, {output.block(), workspace_.block()}); |
| 201 | |
| 202 | if (bias_term_) { |
| 203 | output.device()->Exec([output, this](Context * ctx) { |
| 204 | float beta = 1.f, alpha = 1.0f; |
| 205 | Block *outblock = output.block(), *bblock = this->bias_.block(); |
| 206 | cudnnAddTensor(ctx->cudnn_handle, &alpha, this->bias_desc_, |
| 207 | bblock->data(), &beta, this->y_desc_, |
| 208 | outblock->mutable_data()); |
| 209 | }, {output.block(), bias_.block()}, {output.block()}); |
| 210 | } |
| 211 | return output; |
| 212 | } |
| 213 | |
| 214 | const std::pair<Tensor, vector<Tensor>> CudnnConvolution::Backward( |
| 215 | int flag, const Tensor &grad) { |