chore: import upstream snapshot with attribution

2026-07-13 13:33:03 +08:00
commit 5b57521aa1
8226 changed files with 3425766 additions and 0 deletions
@@ -0,0 +1,139 @@
+//
+//  TFBatchNormalMerge.cpp
+//  MNNConverter
+//
+//  Created by MNN on 2019/09/16.
+//  Copyright © 2018, Alibaba Group Holding Limited
+//
+
+#include "MNN_generated.h"
+#include "TFExtraManager.hpp"
+
+namespace MNN {
+namespace Express {
+class BatchNormalTransform : public TFExtraManager::Transform {
+public:
+    virtual EXPRP onExecute(EXPRP expr) const override {
+        auto inputs                 = expr->inputs();
+        auto op                     = expr->get();
+        std::vector<VARP> subInputs = {inputs[0]};
+        std::unique_ptr<MNN::OpT> BatchNormalOp(new OpT);
+        BatchNormalOp->type       = OpType_BatchNorm;
+        BatchNormalOp->name       = op->name()->str();
+        BatchNormalOp->main.type  = OpParameter_BatchNorm;
+        BatchNormalOp->main.value = new BatchNormT;
+        auto batchnorm            = BatchNormalOp->main.AsBatchNorm();
+        batchnorm->epsilon = 0.001f;
+        bool train = false;
+        auto extra         = op->main_as_Extra();
+        bool nhwc = true;
+        if (nullptr != extra->attr()) {
+            for (int i = 0; i < extra->attr()->size(); ++i) {
+                auto attr = extra->attr()->GetAs<Attribute>(i);
+                if (attr->key()->str() == "epsilon") {
+                    batchnorm->epsilon = attr->f();
+                }
+                if (attr->key()->str() == "is_training") {
+                    train = attr->b();
+                }
+                if (attr->key()->str() == "data_format") {
+                    if (nullptr != attr->s()) {
+                        nhwc = attr->s()->str() == "NHWC";
+                    }
+                }
+            }
+        }
+        auto scaleNode = inputs[1];
+        auto biasNode  = inputs[2];
+        if (train) {
+            std::vector<int> reduceDims = {0, 1, 2};
+            std::vector<int> reshapeDims;
+            if (!nhwc) {
+                // NCHW
+                reduceDims = {0, 2, 3};
+                reshapeDims = {1, -1, 1, 1};
+                scaleNode = _Reshape(scaleNode, reshapeDims);
+                biasNode = _Reshape(biasNode, reshapeDims);
+            }
+            // NHWC, mean for NHW
+            auto mean = _ReduceMean(inputs[0], reduceDims, true);
+            auto xSub = inputs[0] - mean;
+            auto sampleVar      = _ReduceMean(_Square(xSub), reduceDims,
+                                         true); // variance for each channel in the batch
+            auto rSampleStd     = _Reciprocal(_Sqrt(sampleVar + _Const(batchnorm->epsilon)));
+            auto normalizedData = xSub * rSampleStd;
+            auto outputData          = normalizedData * scaleNode + biasNode;
+            outputData->setName(expr->name());
+            return outputData->expr().first;
+        }
+        auto meanNode  = inputs[3];
+        auto varNode   = inputs[4];
+        batchnorm->channels = 0;
+        {
+            auto info = scaleNode->getInfo();
+            auto ptr  = scaleNode->readMap<float>();
+            if (nullptr == info || nullptr == ptr) {
+                MNN_ERROR("Don't support not const scale node \n");
+                return nullptr;
+            }
+            batchnorm->channels = info->size;
+            batchnorm->slopeData.resize(batchnorm->channels);
+            batchnorm->biasData.resize(batchnorm->channels);
+            batchnorm->meanData.resize(batchnorm->channels);
+            batchnorm->varData.resize(batchnorm->channels);
+            ::memcpy(batchnorm->slopeData.data(), ptr, info->size * sizeof(float));
+        }
+        {
+            auto info = biasNode->getInfo();
+            auto ptr  = biasNode->readMap<float>();
+            if (nullptr == info || nullptr == ptr) {
+                MNN_ERROR("Don't support not const bias node \n");
+                return nullptr;
+            }
+            if (info->size != batchnorm->channels) {
+                MNN_ERROR("Don't match channels: %d -> %d\n", batchnorm->channels, info->size);
+                return nullptr;
+            }
+            ::memcpy(batchnorm->biasData.data(), ptr, info->size * sizeof(float));
+        }
+        {
+            auto info = meanNode->getInfo();
+            auto ptr  = meanNode->readMap<float>();
+            if (nullptr == info || nullptr == ptr) {
+                MNN_ERROR("Don't support not const meanNode node \n");
+                return nullptr;
+            }
+            if (info->size != batchnorm->channels) {
+                MNN_ERROR("Don't match channels: %d -> %d\n", batchnorm->channels, info->size);
+                return nullptr;
+            }
+            ::memcpy(batchnorm->meanData.data(), ptr, info->size * sizeof(float));
+        }
+        {
+            auto info = varNode->getInfo();
+            auto ptr  = varNode->readMap<float>();
+            if (nullptr == info || nullptr == ptr) {
+                MNN_ERROR("Don't support not const varNode node \n");
+                return nullptr;
+            }
+            if (info->size != batchnorm->channels) {
+                MNN_ERROR("Don't match channels: %d -> %d\n", batchnorm->channels, info->size);
+                return nullptr;
+            }
+            for (int i = 0; i < batchnorm->varData.size(); ++i) {
+                batchnorm->varData[i] = ptr[i];
+            }
+        }
+        auto newExpr = Expr::create(BatchNormalOp.get(), subInputs, expr->outputSize());
+        return newExpr;
+    }
+};
+static auto gRegister = []() {
+    TFExtraManager::get()->insert("FusedBatchNorm",
+                                  std::shared_ptr<TFExtraManager::Transform>(new BatchNormalTransform));
+    TFExtraManager::get()->insert("FusedBatchNormV3",
+                                  std::shared_ptr<TFExtraManager::Transform>(new BatchNormalTransform));
+    return true;
+}();
+} // namespace Express
+} // namespace MNN