Added axis support for layernorm (#602)

Layernormolization can handle non zero axis now
Added case to verify layernorm with axis 2
Modify layernorm opjson

Type:  Code Improvement

Signed-off-by: Feiyue Chen <Feiyue.Chen@verisilicon.com>
This commit is contained in:
Chen Feiyue 2023-06-15 21:45:46 +08:00 committed by GitHub
parent a64a0f7379
commit fbfbdd7c83
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 58 additions and 9 deletions

View File

@ -3,9 +3,7 @@
"parameters":
[
{"name": "axis",
"dtype": "int32_t",
"Optional": "true",
"default": "0"
"dtype": "int32_t"
},
{"name": "eps",
"dtype": "float",

View File

@ -32,7 +32,7 @@ namespace vx {
namespace ops {
class LayerNormalization : public BuiltinOp {
public:
LayerNormalization(Graph* graph, int32_t axis = 0, float eps = 1e-5f);
LayerNormalization(Graph* graph, int32_t axis, float eps = 1e-5f);
std::shared_ptr<Operation> Clone(std::shared_ptr<Graph>& graph) const override;

View File

@ -32,11 +32,7 @@ namespace vx {
namespace ops {
LayerNormalization::LayerNormalization(Graph* graph, int32_t axis, float eps)
: BuiltinOp(graph, VSI_NN_OP_LAYER_NORM), axis_(axis), eps_(eps) {
// Layer normalization shares the parameters of instance normalization.
if (axis != 0) {
VSILOGE("Layer norm only support axis 0.");
assert(false);
}
this->impl()->node()->nn_param.layernorm.axis = axis_;
this->impl()->node()->nn_param.layernorm.eps = eps_;
}

View File

@ -142,6 +142,61 @@ TEST(LayerNorm, axis_0_shape_2_3_6_1_float) {
EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
}
TEST(LayerNorm, axis_2_shape_4_2_3_1_float) {
auto ctx = tim::vx::Context::Create();
auto graph = ctx->CreateGraph();
tim::vx::ShapeType io_shape({4, 2, 3, 1});
tim::vx::ShapeType param_shape({1,1,3,1});
tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32,
io_shape, tim::vx::TensorAttribute::INPUT);
tim::vx::TensorSpec param_spec(tim::vx::DataType::FLOAT32,
param_shape, tim::vx::TensorAttribute::INPUT);
tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32,
io_shape, tim::vx::TensorAttribute::OUTPUT);
auto input_tensor = graph->CreateTensor(input_spec);
auto gamma_tensor = graph->CreateTensor(param_spec);
auto beta_tensor = graph->CreateTensor(param_spec);
auto output_tensor = graph->CreateTensor(output_spec);
std::vector<float> in_data = {
1, 2, 3, 4,
5, 6, 7, 8,
9, 10, 11, 12,
13, 14, 15, 16,
17, 18, 19, 20,
21, 22, 23, 24};
std::vector<float> gamma = {
1.0f, 1.0f, 1.0f
};
std::vector<float> beta = {
.0f, .0f, .0f
};
std::vector<float> golden = {
-1.22473, -1.22473, -1.22473, -1.22473,
-1.22473, -1.22473, -1.22473, -1.22473,
0, 0, 0, 0,
0, 0, 0, 0,
1.22473, 1.22473, 1.22473, 1.22473,
1.22473, 1.22473, 1.22473, 1.22473
};
EXPECT_TRUE(input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * sizeof(float)));
EXPECT_TRUE(gamma_tensor->CopyDataToTensor(gamma.data(), gamma.size() * sizeof(float)));
EXPECT_TRUE(beta_tensor->CopyDataToTensor(beta.data(), beta.size() * sizeof(float)));
auto op = graph->CreateOperation<tim::vx::ops::LayerNormalization>(2, 0.001);
(*op).BindInputs({input_tensor, beta_tensor, gamma_tensor}).BindOutputs({output_tensor});
EXPECT_TRUE(graph->Compile());
EXPECT_TRUE(graph->Run());
std::vector<float> output(24);
EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
}
#if 0
// Fail case
TEST(LayerNorm, axis_0_shape_3_6_1_uint8) {