Method onForward

tools/train/source/nn/NN.cpp:177–218

    }

    virtual std::vector<Express::VARP> onForward(const std::vector<Express::VARP>& inputs) override {
        Express::VARP x = inputs[0];
        auto dimFormat = x->getInfo()->order;
        VARP outputData = nullptr;
        if (getIsTraining()) {
            if (dimFormat == NC4HW4 || dimFormat == NHWC) {
                x = _Convert(x, NCHW);
            }
            MNN_ASSERT(x->getInfo()->dim[1] == mChannels);
            auto sampleMean     = _ReduceMean(x, mReductionDims, true); // mean for each channel in the batch
            auto xSub = x - sampleMean;
            auto sampleVar      = _ReduceMean(_Square(xSub), mReductionDims,
                                         true); // variance for each channel in the batch
            auto rSampleStd     = _Reciprocal(_Sqrt(sampleVar + _Const(mEps)));
            auto normalizedData = xSub * rSampleStd;
            outputData          = normalizedData * mScale + mBias;

            mRunningMean = _Const(mMomentum) * mRunningMean + _Const(1 - mMomentum) * sampleMean;
            mRunningVariance = _Const(mMomentum) * mRunningVariance + _Const(1 - mMomentum) * sampleVar;
            outputData->setName(name());
            outputData = _Convert(outputData, dimFormat);
            setParameter(mRunningMean, mRunningMeanPos);
            setParameter(mRunningVariance, mRunningVariancePos);
            return {outputData};
        }
        auto rStd  = _Const(1.0f) / _Sqrt(mRunningVariance + _Const(mEps));
        auto alpha = rStd * mScale;
        auto beta  = mBias - mRunningMean * rStd * mScale;
        //outputData = (_Convert(x, NCHW) * alpha) + beta;
        alpha.fix(VARP::CONSTANT);
        beta.fix(VARP::CONSTANT);
        //FUNC_PRINT_ALL(alpha->readMap<float>()[0], f);
        x = _Convert(x, NC4HW4);
        std::vector<float> scale(alpha->getInfo()->size);
        std::vector<float> bias(beta->getInfo()->size);
        ::memcpy(scale.data(), alpha->readMap<float>(), scale.size() * sizeof(float));
        ::memcpy(bias.data(), beta->readMap<float>(), bias.size() * sizeof(float));
        outputData = _Scale(x, mChannels, std::move(scale), std::move(bias));
        outputData->setName(name());
        outputData = _Convert(outputData, dimFormat);
        return {outputData};
    }

private:
    BatchNormModule() = default;