Merge pull request #24 from ROCmSoftwarePlatform/fix-bn-cpu

Fix bn cpu

.gitignore

+*.swp

src/targets/cpu/cpu_lowering.cpp

@@ -24,11 +24,11 @@ T zero(const T&)
 // args[1] -> mini batch mean
 // args[2] -> mini batch variance
 // args[3] -> gamma
-// args[4] -> beta
+// args[4] -> bias
 //
 // The equation to compute batch norm for inference is:
 //
-// output[i] = beta + gamma * (input[i] + mean) / sqrt(variance + epsilon)
+// output[i] = bias + gamma * (input[i] + mean) / sqrt(variance + epsilon)
 //
 // the input data format should be nchw
 //
@@ -46,17 +46,26 @@ struct cpu_batch_norm_inference
 
         double epsilon           = op.epsilon;
         auto input               = args[0];
-        auto mini_batch_mean     = args[1].at<float>();
-        auto mini_batch_variance = args[2].at<float>();
-        auto gamma               = args[3].at<float>();
-        auto beta                = args[4].at<float>();
-
-        visit_all(output, input)([&](auto result, auto buffer) {
-            std::transform(buffer.begin(), buffer.end(), result.begin(), [&](auto x) {
-                return gamma * (x - mini_batch_mean) / std::sqrt(mini_batch_variance + epsilon) +
-                       beta;
+        auto mini_batch_mean     = args[1];
+        auto mini_batch_variance = args[2];
+        auto arg_gamma           = args[3];
+        auto arg_bias            = args[4];
+
+        auto num_batch    = output_shape.lens()[0];
+        auto num_channels = output_shape.lens()[1];
+        auto image_height = output_shape.lens()[2];
+        auto image_width  = output_shape.lens()[3];
+
+        visit_all(output, input, mini_batch_mean, mini_batch_variance, arg_gamma, arg_bias)(
+            [&](auto result, auto buffer, auto mean, auto variance, auto gamma, auto bias) {
+
+                dfor(num_batch, num_channels, image_height, image_width)(
+                    [&](std::size_t n, std::size_t c, std::size_t h, std::size_t w) {
+                        result(n, c, h, w) = gamma(c) * (buffer(n, c, h, w) - mean(c)) /
+                                                 std::sqrt(variance(c) + epsilon) +
+                                             bias(c);
+                    });
             });
-        });
 
         return output;
     }

test/cpu_ops_test.cpp

@@ -9,19 +9,40 @@
 void batch_norm_inference_test()
 {
     migraph::program p;
-    migraph::shape s{migraph::shape::float_type, {4}};
-    auto x        = p.add_literal(migraph::literal{s, {1, 2, 3, 4}});
-    auto gamma    = p.add_literal(migraph::literal{s, {1}});
-    auto beta     = p.add_literal(migraph::literal{s, {0}});
-    auto mean     = p.add_literal(migraph::literal{s, {0}});
-    auto variance = p.add_literal(migraph::literal{s, {1}});
-    p.add_instruction(migraph::batch_norm_inference{}, x, mean, variance, gamma, beta);
+    const size_t width = 2, height = 2, channels = 4, batches = 2;
+    const float x_val = 8.0f, mean_val = 2.0f, variance_val = 4.0f, scale_val = 2.0f,
+                bias_val   = 1.0f;
+    const float output_val = scale_val * (x_val - mean_val) / (std::sqrt(variance_val)) + bias_val;
+
+    migraph::shape s{migraph::shape::float_type, {batches, channels, height, width}};
+    migraph::shape vars{migraph::shape::float_type, {channels}};
+    std::vector<float> x_data(width * height * channels * batches);
+    std::vector<float> scale_data(channels);
+    std::vector<float> bias_data(channels);
+    std::vector<float> mean_data(channels);
+    std::vector<float> variance_data(channels);
+
+    std::fill(x_data.begin(), x_data.end(), x_val);
+    std::fill(mean_data.begin(), mean_data.end(), mean_val);
+    std::fill(variance_data.begin(), variance_data.end(), variance_val);
+    std::fill(scale_data.begin(), scale_data.end(), scale_val);
+    std::fill(bias_data.begin(), bias_data.end(), bias_val);
+
+    auto x        = p.add_literal(migraph::literal{s, x_data});
+    auto scale    = p.add_literal(migraph::literal{vars, scale_data});
+    auto bias     = p.add_literal(migraph::literal{vars, bias_data});
+    auto mean     = p.add_literal(migraph::literal{vars, mean_data});
+    auto variance = p.add_literal(migraph::literal{vars, variance_data});
+
+    p.add_instruction(migraph::batch_norm_inference{}, x, mean, variance, scale, bias);
     p.compile(migraph::cpu::cpu_target{});
     auto result = p.eval({});
-    std::vector<float> result_vector(4);
+
+    std::vector<float> result_vector(width * height * channels * batches);
+    std::vector<float> gold(width * height * channels * batches);
+    std::fill(gold.begin(), gold.end(), output_val);
     result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
-    std::vector<float> gold = {
-        1 / (1 + 1.0e-6), 2 / (1 + 1.0e-6), 3 / (1 + 1.0e-6), 4 / (1 + 1.0e-6)};
+
     EXPECT(test::verify_range(result_vector, gold));
 }