Function apply_on_physical_tensor

imperative/src/impl/ops/batch_norm.cpp:142–202

}

SmallVector<TensorPtr> apply_on_physical_tensor(
        const OpDef& def, const SmallVector<TensorPtr>& inputs,
        SmallVector<LogicalTensorDesc>& output_descs, const bool& validated) {
    auto&& op_def = def.cast_final_safe<BatchNorm>();
    auto&& comp_node = inputs[0]->comp_node();

    DnnOprCaller<megdnn::BN> dnn_opr(comp_node, op_def.param());

    auto src_layout = inputs[0]->layout();
    auto scale_layout = inputs[1]->layout();
    bool empty_input = src_layout.is_empty();
    size_t nr_inp = inputs.size();

    // size_t ws_size = 0, reserve_size = 0;
    size_t reserve_size =
            empty_input ? (size_t)0 : dnn_opr.op()->get_reserve_in_bytes(src_layout);

    // alloc outputs
    auto y = Tensor::make(src_layout, comp_node);
    auto save_mean = Tensor::make(scale_layout, comp_node);
    auto save_variance = Tensor::make(scale_layout, comp_node);
    auto reserve = Tensor::make(TensorLayout{{reserve_size}, dtype::Byte()}, comp_node);

    if (op_def.fwd_mode == ::megdnn::param::BN::FwdMode::INFERENCE) {
        if (!empty_input) {
            dnn_opr.exec_with_ws(
                    inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], save_mean,
                    save_variance, reserve, y);
        }
        return {inputs[3], inputs[4], reserve, y};
    } else {
        if (nr_inp == 5) {
            auto mean = Tensor::make(scale_layout, comp_node);
            auto variance = Tensor::make(scale_layout, comp_node);

            // FIXME
            mean->dev_tensor().copy_from(inputs[3]->dev_tensor());
            variance->dev_tensor().copy_from(inputs[4]->dev_tensor());

            if (!empty_input) {
                dnn_opr.exec_with_ws(
                        inputs[0], inputs[1], inputs[2], mean, variance, save_mean,
                        save_variance, reserve, y);
            }

            return {mean, variance, save_mean, save_variance, reserve, y};
        }

        TensorLayout m_layout({0}, scale_layout.dtype);
        auto mean = Tensor::make(m_layout, comp_node);
        auto variance = Tensor::make(m_layout, comp_node);

        if (!empty_input) {
            dnn_opr.exec_with_ws(
                    inputs[0], inputs[1], inputs[2], mean, variance, save_mean,
                    save_variance, reserve, y);
        }