| 62 | has_init_cudnn_ = true; |
| 63 | } |
| 64 | const Tensor CudnnBatchNorm::Forward(int flag, const Tensor& input) { |
| 65 | auto shape = input.shape(); |
| 66 | auto dtype = input.data_type(); |
| 67 | Tensor output; |
| 68 | Tensor x; |
| 69 | if (is_2d_) |
| 70 | x = Reshape(input, Shape{shape.at(0), shape.at(1), 1, 1}); |
| 71 | else |
| 72 | x = input; |
| 73 | shape = x.shape(); |
| 74 | if (!has_init_cudnn_) { |
| 75 | InitCudnn(shape, dtype); |
| 76 | } else { |
| 77 | int n, c, h, w, s; |
| 78 | cudnnDataType_t type; |
| 79 | CUDNN_CHECK(cudnnGetTensor4dDescriptor(shape_desc_, &type, |
| 80 | &n, &c, &h, &w, &s, &s, &s, &s)); |
| 81 | if (shape[0] != static_cast<size_t>(n)) |
| 82 | InitCudnn(shape, dtype); |
| 83 | CHECK(shape[1] == static_cast<size_t>(c) |
| 84 | && shape[2] == static_cast<size_t>(h) |
| 85 | && shape[3] == static_cast<size_t>(w)) |
| 86 | << "input sample shape should not change" |
| 87 | << "previous shape " << c << ", " << h << ", " << w |
| 88 | << "current shape " << shape[1] << ", " << shape[2] << ", " |
| 89 | << shape[3]; |
| 90 | } |
| 91 | |
| 92 | |
| 93 | // TODO(wangji): check device id of input and params |
| 94 | output.ResetLike(x); |
| 95 | if ((flag & kTrain) == kTrain) { |
| 96 | output.device()->Exec( |
| 97 | [=](Context* ctx) { |
| 98 | Block* inBlock = x.block(), * outBlock = output.block(), |
| 99 | * saveMeanBlock = resultSaveMean_.block(), |
| 100 | * saveVarBlock = resultSaveVariance_.block(), |
| 101 | * runningMeanBlock = runningMean_.block(), |
| 102 | * runningVarBlock = runningVariance_.block(), |
| 103 | * bnScaleBlock = bnScale_.block(), |
| 104 | * bnBiasBlock = bnBias_.block(); |
| 105 | const float alpha = 1.0f, beta = 0.0f; |
| 106 | double epsilon = CUDNN_BN_MIN_EPSILON; |
| 107 | CUDNN_CHECK(cudnnBatchNormalizationForwardTraining( |
| 108 | ctx->cudnn_handle, this->mode_, &alpha, &beta, shape_desc_, |
| 109 | inBlock->data(), shape_desc_, outBlock->mutable_data(), |
| 110 | param_desc_, bnScaleBlock->data(), bnBiasBlock->data(), factor_, |
| 111 | runningMeanBlock->mutable_data(), runningVarBlock->mutable_data(), |
| 112 | epsilon, saveMeanBlock->mutable_data(), |
| 113 | saveVarBlock->mutable_data())); |
| 114 | }, |
| 115 | {x.block(), bnScale_.block(), bnBias_.block()}, |
| 116 | {output.block(), runningMean_.block(), runningVariance_.block(), |
| 117 | resultSaveMean_.block(), resultSaveVariance_.block()}); |
| 118 | buf_.push(x); |
| 119 | } else { |
| 120 | output.device()->Exec( |
| 121 | [=](Context* ctx) { |