chore: import upstream snapshot with attribution

tools/train/source/optimizer/SGD.cpp

@@ -0,0 +1,175 @@
+//
+//  SGD.cpp
+//  MNN
+//
+//  Created by MNN on 2019/11/22.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "SGD.hpp"
+#include "OpGrad.hpp"
+using namespace MNN::Express;
+
+namespace MNN {
+namespace Train {
+SGD::SGD(std::shared_ptr<Module> module) : ParameterOptimizer(module) {
+    auto train = ParameterOptimizer::trainable();
+    for (auto p : train) {
+        mHistory[p] = _Const(0.0f, p->getInfo()->dim, p->getInfo()->order);
+    }
+}
+
+void SGD::setLearningRate(float rate) {
+    mLearningRate = rate;
+}
+
+void SGD::setMomentum(float momentum) {
+    mMomentum = momentum;
+}
+
+void SGD::setWeightDecay(float decay) {
+    mWeightDecay = decay;
+}
+
+void SGD::setRegularizationMethod(RegularizationMethod method) {
+    mRegularizationMethod = method;
+}
+
+float SGD::currentLearningRate() {
+    return mLearningRate;
+}
+
+float SGD::getMomentum() {
+    return mMomentum;
+}
+
+float SGD::getWeightDecay() {
+    return mWeightDecay;
+}
+
+SGD::RegularizationMethod SGD::getRegularizationMethod() {
+    return mRegularizationMethod;
+}
+
+Express::VARP SGD::regularizeParameters(Express::VARP param, Express::VARP grad) {
+    VARP addWeightDecayGrad;
+    if (mRegularizationMethod == L1) {
+        auto temp          = _Sign(param);
+        addWeightDecayGrad = _Const(mWeightDecay, {}, NCHW) * temp + grad;
+    } else if (mRegularizationMethod == L2) {
+        addWeightDecayGrad = _Const(mWeightDecay, {}, NCHW) * param + grad;
+    } else if (mRegularizationMethod == L1L2) {
+        auto temp          = _Sign(param);
+        auto L1 = _Const(mWeightDecay, {}, NCHW) * temp;
+        auto L2 = _Const(mWeightDecay, {}, NCHW) * param;
+        addWeightDecayGrad = L1 + L2 + grad;
+    }
+
+    return addWeightDecayGrad;
+}
+
+Express::VARP SGD::onComputeUpdateValue(Express::VARP param, Express::VARP grad) {
+    auto lr         = _Const(mLearningRate, {}, NCHW);
+    mHistory[param] = lr * grad + _Const(mMomentum, {}, NCHW) * mHistory[param];;
+    mHistory[param].fix(Express::VARP::CONSTANT);
+    //FUNC_PRINT_ALL(_ReduceMax(grad)->readMap<float>()[0], f);
+    return mHistory[param];
+}
+
+std::map<Express::VARP, Express::VARP> SGD::onGetNextParameter(Express::VARP loss) {
+    auto grad = OpGrad::grad(loss, trainable(), mGradBlockExprName);
+    auto parameters = module()->parameters();
+    std::vector<VARP> prepareCompute;
+    for (auto iter : parameters) {
+        if (iter->expr().first->get() != nullptr) {
+            prepareCompute.emplace_back(iter);
+        }
+    }
+    for (auto& iter : grad) {
+        prepareCompute.emplace_back(iter.second);
+    }
+    Variable::prepareCompute(prepareCompute);
+    std::vector<VARP> replaceOp(prepareCompute.size());
+    for (int i=0; i<prepareCompute.size(); ++i) {
+        auto info = prepareCompute[i]->getInfo();
+        auto ptr = prepareCompute[i]->readMap<void>();
+        if (nullptr == ptr) {
+            MNN_ERROR("Compute error in SGD\n");
+            return {};
+        }
+        auto newVar = _Const(ptr, info->dim, info->order, info->type);
+        replaceOp[i]= newVar;
+    }
+    for (int i=0; i<prepareCompute.size(); ++i) {
+        Variable::replace(prepareCompute[i], replaceOp[i]);
+    }
+
+    for (auto& iter : grad) {
+        // apply regularization
+        auto addWeightDecayGrad = regularizeParameters(iter.first, iter.second);
+        addWeightDecayGrad.fix(Express::VARP::CONSTANT);
+        // apply momentum, etc.
+        auto updateValue = this->onComputeUpdateValue(iter.first, addWeightDecayGrad);
+        // apply update
+        auto newParameter = iter.first - updateValue;
+        iter.second       = newParameter;
+    }
+    return grad;
+}
+
+std::pair<std::vector<Express::VARP>, std::vector<Express::VARP>>  SGD::onMakeParameterUpdateGraphByGrad(const std::vector<ParameterOptGrad>& parameterGrads) {
+    std::map<MNN::Express::VARP, MNN::Express::VARP> varUpdateMap;
+    auto momentum = _Const(mMomentum, {}, NCHW);
+    auto weightDecay = _Const(mWeightDecay, {}, NCHW);
+    for (int pIndex=0; pIndex<parameterGrads.size(); ++pIndex) {
+        auto p = parameterGrads[pIndex].parameter;
+        auto grad = parameterGrads[pIndex].parameterGrad;
+
+        // FIXME
+#if 0
+        if (p->name().find("_BN_RunningMean_Weight") != string::npos) {
+            varUpdateMap[p] = trainInfo[p->name()];
+            continue; // not update running stats
+        }
+        if (p->name().find("_BN_RunningVariance_Weight") != string::npos) {
+            varUpdateMap[p] = trainInfo[p->name()];
+            continue; // not update running stats
+        }
+        if (p->name().find("_BN_Eps_Weight") != string::npos) {
+            continue; // not update eps
+        }
+#endif
+        auto pInfo = p->getInfo();
+        auto pDims = pInfo->dim;
+
+        auto l2grad = weightDecay * p;
+        l2grad->setName(p->name() + "_l2grad");
+
+        VARP gradWithDecay = grad + l2grad;
+
+        VARP history = _Const(0.0f, pDims, pInfo->order);
+        history->setName(p->name() + "_momentum");
+        history.fix(VARP::TRAINABLE);
+
+        auto newHistory = gradWithDecay + momentum * history;
+        newHistory->setName("update_" + history->name());
+
+        VARP localLearningRate = parameterGrads[pIndex].learningRate;
+        VARP finalGrad = localLearningRate * history;
+        finalGrad->setName(p->name() + "_final_grad");
+
+        auto updateValue = _Subtract(p, finalGrad);
+        updateValue->setName("update_" + p->name());
+        varUpdateMap[p] = updateValue;
+        varUpdateMap[history] = newHistory;
+    }
+    std::vector<Express::VARP> res;
+    std::vector<Express::VARP> resUpdate;
+    for (auto& iter : varUpdateMap) {
+        res.push_back(iter.first);
+        resUpdate.push_back(iter.second);
+    }
+    return std::make_pair(res, resUpdate);
+}
+} // namespace Train
+} // namespace MNN