batch_norm.cpp 2.7 KB
Newer Older
1 2 3 4
/**
 * \file imperative/src/impl/ops/batch_norm.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
5
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
6 7 8 9 10 11
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

12 13
#include "megbrain/imperative/ops/autogen.h"
#include "megbrain/opr/dnn/batch_norm.h"
14 15 16 17 18 19 20 21 22
#include "../op_trait.h"

namespace mgb {
namespace imperative {

namespace {

std::shared_ptr<OpDef> make_from_op_node(cg::OperatorNodeBase* node_) {
    auto* node = &node_->cast_final_safe<opr::BatchNorm>();
23
    return BatchNorm::make(node->param());
24 25 26 27 28 29 30 31 32
}

cg::OperatorNodeBase* apply_on_var_node(
        const OpDef& def,
        const VarNodeArray& inputs) {
    auto&& bn_opr = def.cast_final_safe<BatchNorm>();
    size_t nr_inp = inputs.size();
    mgb_assert(nr_inp == 3 ||nr_inp == 5,
              "BatchNorm expects 3 or 5 inputs; got %lu actually", nr_inp);
33
    OperatorNodeConfig config{bn_opr.make_name()};
34 35
    if (nr_inp == 3) {
        return opr::BatchNorm::make(
36
            inputs[0], inputs[1], inputs[2], bn_opr.param(), config)[0]
37 38 39
            .node()->owner_opr();
    } else {
        return opr::BatchNorm::make(
40
            inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], bn_opr.param(), config)[0]
41 42 43 44
            .node()->owner_opr();
    }
}

45
std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible(
46 47 48 49 50 51 52
        const OpDef& def,
        const SmallVector<LogicalTensorDesc>& inputs) {
    auto&& op_def = def.cast_final_safe<BatchNorm>();
    size_t nr_inp = inputs.size();
    mgb_assert(nr_inp == 3 ||nr_inp == 5,
              "BatchNorm expects 3 or 5 inputs; got %lu actually", nr_inp);
    // need running mean/variance
53
    bool need_stat = (nr_inp == 5) && op_def.fwd_mode == BatchNorm::FwdMode::TRAINING;
54 55 56 57
    size_t nr_out = need_stat? 5 : 3;
    SmallVector<LogicalTensorDesc> out_shapes(nr_out);
    auto&& i0 = inputs[0];
    auto&& i1 = inputs[1];
58 59
    // [running_mean, running_var,] save_mean, save_var
    for (size_t i = 0; i < nr_out-1; ++ i) {
60 61
        out_shapes[i] = {i1.layout, i1.comp_node};
    }
62
    // output tensor
63
    out_shapes[nr_out-1] = {i0.layout, i0.comp_node};
64
    return {out_shapes, out_shapes[nr_out-1].layout.ndim != 0};
65 66 67 68 69 70 71 72 73 74 75 76 77
}

OP_TRAIT_REG(BatchNorm, BatchNorm, opr::BatchNorm)
    .make_from_op_node(make_from_op_node)
    .apply_on_var_node(apply_on_var_node)
    .infer_output_attrs_fallible(infer_output_attrs_fallible)
    .fallback();
} // anonymous namespace

}  // namespace imperative
}  // namespace mgb

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}