| 78 | class BatchNormModule : public Module { |
| 79 | public: |
| 80 | BatchNormModule(EXPRP expr, const float m = 0.99) { |
| 81 | MNN_ASSERT(expr->get() != nullptr); |
| 82 | MNN_ASSERT(expr->get()->type() == OpType_BatchNorm); |
| 83 | auto bnPa = expr->get()->main_as_BatchNorm(); |
| 84 | auto& inputs = expr->inputs(); |
| 85 | int dims = 4; |
| 86 | if (!inputs.empty()) { |
| 87 | auto info = inputs[0]->getInfo(); |
| 88 | if (nullptr != info) { |
| 89 | dims = info->dim.size(); |
| 90 | } |
| 91 | } |
| 92 | mEps = bnPa->epsilon(); |
| 93 | mMomentum = m; |
| 94 | mChannels = bnPa->channels(); |
| 95 | std::vector<int> statShape; |
| 96 | std::vector<int> reductionDims; |
| 97 | int channels = mChannels; |
| 98 | if (dims == 2) { |
| 99 | statShape = {1, channels}; |
| 100 | mReductionDims = {0}; |
| 101 | } |
| 102 | if (dims == 3) { |
| 103 | statShape = {1, channels, 1}; |
| 104 | mReductionDims = {0, 2}; |
| 105 | } |
| 106 | if (dims == 4) { |
| 107 | statShape = {1, channels, 1, 1}; |
| 108 | mReductionDims = {0, 2, 3}; |
| 109 | } |
| 110 | MNN_ASSERT(bnPa->biasData()->size() == mChannels); |
| 111 | mBias = _TrainableParam(bnPa->biasData()->data(), statShape, NCHW); |
| 112 | MNN_ASSERT(bnPa->slopeData()->size() == mChannels); |
| 113 | mScale = _TrainableParam(bnPa->slopeData()->data(), statShape, NCHW); |
| 114 | MNN_ASSERT(bnPa->meanData()->size() == mChannels); |
| 115 | mRunningMean = _Const(bnPa->meanData()->data(), statShape, NCHW); |
| 116 | MNN_ASSERT(bnPa->meanData()->size() == mChannels); |
| 117 | mRunningVariance = _Const(bnPa->varData()->data(), statShape, NCHW); |
| 118 | addParameter(mScale); |
| 119 | addParameter(mBias); |
| 120 | mRunningVariancePos = addParameter(mRunningVariance); |
| 121 | mRunningMeanPos = addParameter(mRunningMean); |
| 122 | |
| 123 | setType("BatchNorm"); |
| 124 | } |
| 125 | BatchNormModule(const int channels, const int dims = 4, const float m = 0.99, const float e = 1e-5) { |
| 126 | mMomentum = m; |
| 127 | mEps = e; |