#include "megbrain/opr/dnn/layer_norm.h"
#include "megbrain/comp_node_env.h"
#include "megbrain/test/autocheck.h"
#include "megbrain/test/helper.h"
#include "megbrain/test/megdnn_helper.h"
#include "megdnn/oprs.h"
#include <cmath>
#include <iomanip>
#include <random>
#include <sstream>
using namespace mgb;
namespace {
using Param = opr::LayerNormForward::Param;
void run_forward(bool is_affine, size_t normalized_size) {
using Checker = AutoOprChecker<3, 3>;
Param param;
param.eps = 1e-5;
param.affine = is_affine;
param.normalized_dim = 1;
param.normalized_size = normalized_size;
auto make_graph = [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {
auto out = opr::LayerNormForward::make(inputs[0], inputs[1], inputs[2], param);
return {out[0], out[1], out[2]};
};
auto fwd = [&](Checker::NumOutArray& dest, Checker::NumInpArray inp) {
auto opr =
MegDNNHandle::get(CompNodeEnv::from_comp_node(CompNode::default_cpu()))
->create_operator<megdnn::LayerNormForward>();
auto inp_shape = inp[0]->shape();
auto n_slices = inp_shape[0];
auto slice_len = inp_shape[1];
opr->param() = param;
dest[0].dtype(dtype::Float32())
.comp_node(inp[0]->comp_node())
.resize({n_slices, slice_len});
dest[1].dtype(dtype::Float32())
.comp_node(inp[0]->comp_node())
.resize({n_slices});
dest[2].dtype(dtype::Float32())
.comp_node(inp[0]->comp_node())
.resize({n_slices});
opr->exec(
inp[0]->as_megdnn(), inp[1]->as_megdnn(), inp[2]->as_megdnn(),
dest[0].as_megdnn(), dest[1].as_megdnn(), dest[2].as_megdnn(), {});
};
auto gen = [&](HostTensorND& src) {
HostTensorGenerator<dtype::Float32, RandomDistribution::GAUSSIAN> src_gen(0.f);
src = *src_gen(src.shape(), src.comp_node());
};
Checker::RunOptions option;
option.numdiff_max_err = 1e-4;
Checker checker{make_graph, fwd};
checker.set_input_generator(0, gen);
checker.set_input_generator(1, gen);
checker.set_input_generator(2, gen);
checker.set_input_allow_grad(0, false);
checker.set_input_allow_grad(1, false);
checker.set_input_allow_grad(2, false);
checker.set_output_allow_grad(0, false);
checker.set_output_allow_grad(1, false);
checker.set_output_allow_grad(2, false);
checker.run({TensorShape{normalized_size, normalized_size},
TensorShape{normalized_size}, TensorShape{normalized_size}},
option)
.run({TensorShape{normalized_size, normalized_size},
TensorShape{normalized_size}, TensorShape{normalized_size}},
option)
.run({TensorShape{normalized_size, normalized_size},
TensorShape{normalized_size}, TensorShape{normalized_size}},
option);
}
TEST(TestOprDNN, LayerNormForwardAffine) {
REQUIRE_GPU(1);
run_forward(true, 1);
run_forward(true, 16);
run_forward(true, 17);
}
}