Add the BatchNorm train/test flag support

include/aidge/operator/BatchNorm.hpp

@@ -24,7 +24,7 @@
 
 namespace Aidge {
 
-enum class BatchNormAttr { Epsilon, Momentum };
+enum class BatchNormAttr { Epsilon, Momentum, TrainingMode };
 
 template <DimIdx_t DIM>
 class BatchNorm_Op : public OperatorTensor,
@@ -33,7 +33,7 @@ public:
     static const std::string Type;
 
 private:
-    using Attributes_ = StaticAttributes<BatchNormAttr, float, float>;
+    using Attributes_ = StaticAttributes<BatchNormAttr, float, float, bool>;
     template <BatchNormAttr e>
     using attr = typename Attributes_::template attr<e>;
     const std::shared_ptr<Attributes_> mAttributes;
@@ -42,7 +42,7 @@ public:
 
     BatchNorm_Op() = delete;
 
-    constexpr BatchNorm_Op(float epsilon, float momentum)
+    constexpr BatchNorm_Op(float epsilon, float momentum, bool trainingMode)
         : OperatorTensor(Type,
                             {InputCategory::Data,
                                 InputCategory::Param,
@@ -52,7 +52,9 @@ public:
                             1),
           mAttributes(std::make_shared<Attributes_>(
             attr<BatchNormAttr::Epsilon>(epsilon),
-            attr<BatchNormAttr::Momentum>(momentum))) {}
+            attr<BatchNormAttr::Momentum>(momentum),
+            attr<BatchNormAttr::TrainingMode>(trainingMode)
+            )) {}
 
     /**
      * @brief Copy-constructor. Copy the operator attributes and its output tensor(s), but not its input tensors (the new operator has no input associated).
@@ -84,6 +86,7 @@ public:
     inline std::shared_ptr<Attributes> attributes() const override { return mAttributes; }
     inline float& epsilon() const { return mAttributes->template getAttr<BatchNormAttr::Epsilon>(); }
     inline float& momentum() const { return mAttributes->template getAttr<BatchNormAttr::Momentum>(); }
+    inline bool& trainingMode() const { return mAttributes->template getAttr<BatchNormAttr::TrainingMode>(); }
 
     static const std::vector<std::string> getInputsName() {
         return {"data_input", "scale", "shift", "mean", "variance"};
@@ -101,16 +104,17 @@ template <DimSize_t DIM>
 std::shared_ptr<Node> BatchNorm(const DimSize_t nbFeatures,
                                        const float epsilon = 1.0e-5F,
                                        const float momentum = 0.1F,
+                                       const bool trainingMode = false,
                                        const std::string& name = "");
 }  // namespace Aidge
 
-extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<2>(const DimSize_t, const float, const float, const std::string&);
-extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<3>(const DimSize_t, const float, const float, const std::string&);
-extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<4>(const DimSize_t, const float, const float, const std::string&);
+extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<2>(const DimSize_t, const float, const float, const bool, const std::string&);
+extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<3>(const DimSize_t, const float, const float, const bool, const std::string&);
+extern template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<4>(const DimSize_t, const float, const float, const bool, const std::string&);
 
 namespace {
 template <>
-const char *const EnumStrings<Aidge::BatchNormAttr>::data[] = { "epsilon", "momentum" };
+const char *const EnumStrings<Aidge::BatchNormAttr>::data[] = { "epsilon", "momentum", "training_mode" };
 }
 
 #endif //AIDGE_CORE_OPERATOR_BATCHNORM_H_

python_binding/operator/pybind_BatchNorm.cpp

@@ -26,16 +26,17 @@ void declare_BatchNormOp(py::module& m) {
     const std::string pyClassName("BatchNorm" + std::to_string(DIM) + "DOp");
     py::class_<BatchNorm_Op<DIM>, std::shared_ptr<BatchNorm_Op<DIM>>, OperatorTensor>(
     m, pyClassName.c_str(), py::multiple_inheritance())
-        .def(py::init<float, float>(),
+        .def(py::init<float, float, bool>(),
             py::arg("epsilon"),
-            py::arg("momentum"))
+            py::arg("momentum"),
+            py::arg("training_mode"))
         .def_static("get_inputs_name", &BatchNorm_Op<DIM>::getInputsName)
         .def_static("get_outputs_name", &BatchNorm_Op<DIM>::getOutputsName)
         .def_readonly_static("Type", &BatchNorm_Op<DIM>::Type);
 
     declare_registrable<BatchNorm_Op<DIM>>(m, pyClassName);
 
-    m.def(("BatchNorm" + std::to_string(DIM) + "D").c_str(), &BatchNorm<DIM>, py::arg("nb_features"), py::arg("epsilon") = 1.0e-5F, py::arg("momentum") = 0.1F, py::arg("name") = "");
+    m.def(("BatchNorm" + std::to_string(DIM) + "D").c_str(), &BatchNorm<DIM>, py::arg("nb_features"), py::arg("epsilon") = 1.0e-5F, py::arg("momentum") = 0.1F, py::arg("training_mode") = false, py::arg("name") = "");
 }
 
 void init_BatchNorm(py::module &m) {

src/operator/BatchNorm.cpp

@@ -108,9 +108,10 @@ template <Aidge::DimSize_t DIM>
 inline std::shared_ptr<Aidge::Node> Aidge::BatchNorm(const Aidge::DimSize_t nbFeatures,
                                        const float epsilon,
                                        const float momentum,
+                                       const bool trainingMode,
                                        const std::string& name) {
     static_assert(DIM<=MaxDim,"Too many kernel dimensions required by BatchNorm, not supported");
-    auto batchNorm = std::make_shared<Node>(std::make_shared<BatchNorm_Op<static_cast<DimIdx_t>(DIM)>>(epsilon, momentum), name);
+    auto batchNorm = std::make_shared<Node>(std::make_shared<BatchNorm_Op<static_cast<DimIdx_t>(DIM)>>(epsilon, momentum, trainingMode), name);
     addProducer(batchNorm, 1, {nbFeatures}, "scale");
     addProducer(batchNorm, 2, {nbFeatures}, "shift");
     addProducer(batchNorm, 3, {nbFeatures}, "batch_mean");
@@ -118,6 +119,6 @@ inline std::shared_ptr<Aidge::Node> Aidge::BatchNorm(const Aidge::DimSize_t nbFe
     return batchNorm;
 }
 
-template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<2>(const DimSize_t, const float, const float, const std::string&);
-template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<3>(const DimSize_t, const float, const float, const std::string&);
-template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<4>(const DimSize_t, const float, const float, const std::string&);
+template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<2>(const DimSize_t, const float, const float, const bool, const std::string&);
+template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<3>(const DimSize_t, const float, const float, const bool, const std::string&);
+template std::shared_ptr<Aidge::Node> Aidge::BatchNorm<4>(const DimSize_t, const float, const float, const bool, const std::string&);

unit_tests/graph/Test_Matching.cpp

@@ -352,11 +352,11 @@ TEST_CASE("[core/graph] Matching") {
     auto g2 = Sequential({
         Producer({16, 3, 512, 512}, "dataProvider"),
         Conv(3, 4, {5, 5}, "conv1"),
-        BatchNorm<2>(4, 1.0e-5, 0.1, "bn1"),
+        BatchNorm<2>(4, 1.0e-5, 0.1, false, "bn1"),
         Conv(4, 4, {5, 5}, "conv2"),
         ReLU("relu2"),
         Conv(4, 4, {5, 5}, "conv3"),
-        BatchNorm<2>(4, 1.0e-5, 0.1, "bn3"),
+        BatchNorm<2>(4, 1.0e-5, 0.1, false, "bn3"),
         FC(4, 4, false, "fc1"),
         FC(4, 4, false, "fc2"),
         FC(4, 4, false, "fc3"),