Merge branch 'dev' into 'main'

version 0.2.0

See merge request eclipse/aidge/aidge_quantization!15

include/aidge/operator/SAT/DoReFa.hpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_CORE_OPERATOR_DOREFA_H_
+#define AIDGE_CORE_OPERATOR_DOREFA_H_
+
+#include <cassert>
+#include <memory>
+#include <vector>
+
+#include "aidge/backend/OperatorImpl.hpp"
+#include "aidge/graph/Node.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+#include "aidge/utils/ErrorHandling.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/StaticAttributes.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+
+enum class DoReFaAttr { Range, Mode };
+enum class DoReFaMode {
+    Default,        // Original SAT paper (not including 0)
+    Symmetric,      // Symmetric range including 0
+    //Asymmetric,
+    //FullRange
+};
+
+/**
+ * DoReFa is the weights quantizer for the 2nd training phase (quantization) of the SAT method.
+ */
+class DoReFa_Op : public OperatorTensor,
+    public Registrable<DoReFa_Op, std::string, std::function<std::shared_ptr<OperatorImpl>(const DoReFa_Op&)>> {
+public:
+    static const std::string Type;
+
+private:
+    using Attributes_ = StaticAttributes<DoReFaAttr, size_t, DoReFaMode>;
+    template <DoReFaAttr e> using attr = typename Attributes_::template attr<e>;
+    const std::shared_ptr<Attributes_> mAttributes;
+
+public:
+    DoReFa_Op(size_t range = 255, DoReFaMode mode = DoReFaMode::Default)
+      : OperatorTensor(Type, {InputCategory::Param}, 1),
+        mAttributes(std::make_shared<Attributes_>(
+            attr<DoReFaAttr::Range>(range),
+            attr<DoReFaAttr::Mode>(mode)))
+    {}
+
+    /**
+     * @brief Copy-constructor. Copy the operator attributes and its output tensor(s), but not its input tensors (the new operator has no input associated).
+     * @param op Operator to copy.
+     */
+    DoReFa_Op(const DoReFa_Op& op)
+        : OperatorTensor(op),
+          mAttributes(op.mAttributes)
+    {
+        if (op.mImpl){
+            SET_IMPL_MACRO(DoReFa_Op, *this, op.backend());
+        }else{
+            mImpl = nullptr;
+        }
+    }
+
+    /**
+     * @brief Clone the operator using its copy-constructor.
+     * @see Operator::DoReFa_Op
+     */
+    std::shared_ptr<Operator> clone() const override {
+        return std::make_shared<DoReFa_Op>(*this);
+    }
+
+    std::set<std::string> getAvailableBackends() const override final;
+    void setBackend(const std::string& name, DeviceIdx_t device = 0) override final;
+
+    inline std::shared_ptr<Attributes> attributes() const override { return mAttributes; }
+    inline size_t& range() const noexcept { return mAttributes->getAttr<DoReFaAttr::Range>(); }
+    inline DoReFaMode& mode() const noexcept { return mAttributes->getAttr<DoReFaAttr::Mode>(); }
+
+    static const std::vector<std::string> getInputsName(){
+        return {"data_input"};
+    }
+    static const std::vector<std::string> getOutputsName(){
+        return {"data_output"};
+    }
+};
+
+inline std::shared_ptr<Node> DoReFa(size_t range = 255, DoReFaMode mode = DoReFaMode::Default, const std::string& name = "") {
+    return std::make_shared<Node>(std::make_shared<DoReFa_Op>(range, mode), name);
+}
+}
+
+namespace {
+template <>
+const char *const EnumStrings<Aidge::DoReFaAttr>::data[] = {"range", "mode"};
+
+template <>
+const char *const EnumStrings<Aidge::DoReFaMode>::data[] = {"default", "symmetric", "asymmetric", "full_range"};
+}
+
+#endif /* AIDGE_CORE_OPERATOR_DOREFA_H_ */

include/aidge/operator/SAT/TanhClamp.hpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_CORE_OPERATOR_TANHCLAMP_H_
+#define AIDGE_CORE_OPERATOR_TANHCLAMP_H_
+
+#include <cassert>
+#include <memory>
+#include <vector>
+
+#include "aidge/backend/OperatorImpl.hpp"
+#include "aidge/graph/Node.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+#include "aidge/operator/Producer.hpp"
+#include "aidge/utils/ErrorHandling.hpp"
+#include "aidge/utils/Registrar.hpp"
+#include "aidge/utils/StaticAttributes.hpp"
+#include "aidge/utils/Types.h"
+
+namespace Aidge {
+
+/**
+ * TanhClamp is the weights clamping for the 1st training phase (clamping) of the SAT method.
+ */
+class TanhClamp_Op : public OperatorTensor,
+    public Registrable<TanhClamp_Op, std::string, std::function<std::shared_ptr<OperatorImpl>(const TanhClamp_Op&)>> {
+
+public:
+    static const std::string Type;
+
+    TanhClamp_Op()
+      : OperatorTensor(Type, {InputCategory::Data}, 2)
+    {}
+
+    /**
+     * @brief Copy-constructor. Copy the operator attributes and its output tensor(s), but not its input tensors (the new operator has no input associated).
+     * @param op Operator to copy.
+     */
+    TanhClamp_Op(const TanhClamp_Op& op)
+        : OperatorTensor(op)
+    {
+        if (op.mImpl){
+            SET_IMPL_MACRO(TanhClamp_Op, *this, op.backend());
+        }else{
+            mImpl = nullptr;
+        }
+    }
+
+    /**
+     * @brief Clone the operator using its copy-constructor.
+     * @see Operator::TanhClamp_Op
+     */
+    std::shared_ptr<Operator> clone() const override {
+        return std::make_shared<TanhClamp_Op>(*this);
+    }
+
+    bool forwardDims(bool allowDataDependency = false) override final;
+    std::set<std::string> getAvailableBackends() const override final;
+
+    void setBackend(const std::string& name, DeviceIdx_t device = 0) override final;
+
+    static const std::vector<std::string> getInputsName(){
+        return {"data_input"};
+    }
+    static const std::vector<std::string> getOutputsName(){
+        return {"data_output", "scaling"};
+    }
+};
+
+inline std::shared_ptr<Node> TanhClamp(const std::string& name = "") {
+    return std::make_shared<Node>(std::make_shared<TanhClamp_Op>(), name);
+}
+}
+
+#endif /* AIDGE_CORE_OPERATOR_TANHCLAMP_H_ */

include/aidge/quantization/PTQ/Clip.hpp

@@ -26,7 +26,7 @@ namespace Aidge
     /**
      * @brief Kind of clipping policy to apply during the activation quantization
      */
-    enum Clipping {MAX = 1, MSE, AA, KL};
+    enum class Clipping {MAX = 1, MSE, AA, KL};
 
     /**
      * @brief Compute the histograms of the activations of each node contained in the map of the ranges (passed as argument).

include/aidge/quantization/PTQ/PTQ.hpp

@@ -26,12 +26,12 @@ namespace Aidge {
     /**
      * @brief Set of the types of the nodes which contain affine transforms (that is Y = A.X + B)
      */
-    static const std::set<std::string> affineNodeTypes({"FC", "Conv", "ConvDepthWise", "PaddedConv", "PaddedConvDepthWise"});
+    static const std::set<std::string> affineNodeTypes({"FC", "Conv2D", "ConvDepthWise2D", "PaddedConv2D", "PaddedConvDepthWise2D"});
 
     /**
      * @brief Set of the types of the nodes which does not affect the PTQ process
      */
-    static const std::set<std::string> seamlessNodeTypes({"Pad", "MaxPooling", "AvgPooling", "PaddedMaxPooling", "PaddedAvgPooling", "GlobalAveragePooling", "Reshape", "Transpose", "Gather"});
+    static const std::set<std::string> seamlessNodeTypes({"Pad2D", "MaxPooling2D", "AvgPooling2D", "PaddedMaxPooling2D", "PaddedAvgPooling2D", "GlobalAveragePooling", "Reshape", "Transpose", "Gather"});
 
     /**
      * @brief Set of the types of the nodes that merge multiple branches into one
@@ -74,6 +74,10 @@ namespace Aidge {
      */
     bool checkArchitecture(std::shared_ptr<GraphView> graphView);
 
+
+    void prepareNetwork(std::shared_ptr<GraphView> graphView);
+
+
     /**
      * @brief Insert a scaling node after each affine node of the GraphView.
      * Also insert a scaling node in every purely residual branches.
@@ -154,7 +158,6 @@ namespace Aidge {
      * @param graphView The GraphView under test.
      */
     void devPTQ(std::shared_ptr<GraphView> graphView); 
-
 }
 
 #endif /* AIDGE_QUANTIZATION_PTQ_PTQ_H_ */

include/aidge/quantization/QAT/QAT_FixedQ.hpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_QUANTIZATION_QAT_FIXEDQ_H_
+#define AIDGE_QUANTIZATION_QAT_FIXEDQ_H_
+
+#include "aidge/graph/Node.hpp"
+#include "aidge/graph/GraphView.hpp"
+#include "aidge/data/Tensor.hpp"
+
+namespace Aidge {
+namespace QuantFixedQ {
+
+/**
+ * @brief Insert the FixedQ quantizer nodes in a given GraphView
+ * @param graphView The GraphView containing the graph to quantize.
+ * @param nbBits Number of quantization bits.
+ * @param span Fixed output span of the quantizers.
+ */
+void insertQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, float span);
+
+/**
+ * @brief Given a GraphView with parameters properly initialized and some calibration data,
+ * insert the FixedQ quantizer nodes, and adjust their output spans.
+ * @param graphView The GraphView containing the graph to quantize.
+ * @param nbBits Number of quantization bits.
+ * @param calibrationData Calibration data used to adjust the spans.
+ * @param scale Multiplicative constant applied to the spans.
+ */
+void insertAndInitQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, std::shared_ptr<Tensor> calibrationData, float scale);
+
+/**
+ * @brief Developement and test routine.
+ * @param graphView The GraphView under test.
+ */
+void devQAT(std::shared_ptr<GraphView> graphView); 
+
+}
+}
+
+#endif /* AIDGE_QUANTIZATION_QAT_FIXEDQ_H_ */
+

include/aidge/quantization/QAT/QAT_LSQ.hpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#ifndef AIDGE_QUANTIZATION_QAT_LSQ_H_
+#define AIDGE_QUANTIZATION_QAT_LSQ_H_
+
+#include "aidge/graph/Node.hpp"
+#include "aidge/graph/GraphView.hpp"
+#include "aidge/data/Tensor.hpp"
+
+namespace Aidge {
+namespace QuantLSQ {
+
+/**
+ * @brief Insert the LSQ quantizer nodes in a given GraphView
+ * @param graphView The GraphView containing the graph to quantize.
+ * @param nbBits Number of quantization bits.
+ * @param span Fixed output span of the quantizers.
+ */
+void insertQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, float step_size);
+
+/**
+ * @brief Given a GraphView with parameters properly initialized and some calibration data,
+ * insert the LSQ quantizer nodes, and adjust their step-sizes.
+ * @param graphView The GraphView containing the graph to quantize.
+ * @param nbBits Number of quantization bits.
+ * @param calibrationData Calibration data used to adjust the spans.
+ * @param scale Multiplicative constant applied to the spans.
+ */
+void insertAndInitQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, std::shared_ptr<Tensor> calibrationData);
+
+}
+}
+
+#endif /* AIDGE_QUANTIZATION_QAT_LSQ_H_ */
+

pyproject.toml

+[project]
+name = "aidge_quantization"
+description="Quantization algorithms to compress aidge networks."
+dependencies = [
+    "numpy>=1.21.6",
+]
+requires-python = ">= 3.7"
+readme = "README.md"
+license = { file = "LICENSE" }
+classifiers = [ 
+    "Development Status :: 2 - Pre-Alpha",
+    "Programming Language :: Python :: 3"
+    ]
+dynamic = ["version"] # defined in tool.setuptools_scm
+# version="1"
+
+[build-system]
+requires = [
+    "setuptools>=64",
+    "setuptools_scm[toml]==7.1.0",
+    "cmake>=3.15.3.post1",
+    "toml"
+]
+build-backend = "setuptools.build_meta"
+
+#####################################################
+# SETUPTOOLS
+[tool.setuptools]
+include-package-data = true
+[tool.setuptools.packages.find]
+where = ["."]  # list of folders that contain the packages (["."] by default)
+include = ["aidge_quantization*"]  # package names should match these glob patterns (["*"] by default)
+exclude = ["aidge_quantization.unit_tests*"]  # exclude packages matching these glob patterns (empty by default)
+
+# SETUPTOOLS_SCM
+[tool.setuptools_scm]
+write_to = "aidge_quantization/_version.py"
+
+#####################################################
+# CIBUILDWHEEL
+[tool.cibuildwheel]
+build-frontend = "build"
+test-requires = "pytest"
+test-command = "pytest {package}/aidge_quantization/unit_tests"
+# uncomment to run cibuildwheel locally on selected distros
+# build=[
+# "cp38-manylinux_x86_64",
+# "cp39-manylinux_x86_64",
+# "cp310-manylinux_x86_64",
+# "cp38-win_amd64",
+# "cp39-win_amd64",
+# "cp310-win_amd64",
+# ]
+## AIDGE DEPENDENCIES DECLARATION
+[tool.cibuildwheel.environment]
+AIDGE_DEPENDENCIES = "aidge_core aidge_backend_cpu aidge_onnx" # format => "dep_1 dep_2 ... dep_n"
+AIDGE_INSTALL="/AIDGE_INSTALL_CIBUILDWHEEL"
+[tool.cibuildwheel.linux]
+before-test = [
+    "bash .gitlab/ci/cibuildwheel_build_deps_before_build_wheel.sh /host"
+]
+before-build = [
+    "bash .gitlab/ci/cibuildwheel_build_deps_before_build_wheel.sh /host"
+]
+[tool.cibuildwheel.windows]
+before-build = [
+    "powershell -File .\\.gitlab\\ci\\cibuildwheel_build_deps_before_build_wheel.ps1"
+]
+before-test = [
+    "powershell -File .\\.gitlab\\ci\\cibuildwheel_build_deps_before_build_wheel.ps1"
+]

python_binding/operator/SAT/pybind_DoReFa.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/SAT/DoReFa.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+
+namespace py = pybind11;
+namespace Aidge {
+
+void init_DoReFa(py::module& m) {
+    py::enum_<DoReFaMode>(m, "DoReFaMode")
+        .value("Default", DoReFaMode::Default)
+        .value("Symmetric", DoReFaMode::Symmetric)
+        .export_values();
+
+    py::class_<DoReFa_Op, std::shared_ptr<DoReFa_Op>, OperatorTensor>(m, "DoReFaOp", py::multiple_inheritance())
+    .def(py::init<size_t, DoReFaMode>(), py::arg("range") = 255, py::arg("mode") = DoReFaMode::Default)
+    .def_static("get_inputs_name", &DoReFa_Op::getInputsName)
+    .def_static("get_outputs_name", &DoReFa_Op::getOutputsName);
+    declare_registrable<DoReFa_Op>(m, "DoReFaOp");
+    m.def("DoReFa", &DoReFa, py::arg("range") = 255, py::arg("mode") = DoReFaMode::Default, py::arg("name") = "");
+}
+}  // namespace Aidge

python_binding/operator/SAT/pybind_TanhClamp.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/SAT/TanhClamp.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+
+namespace py = pybind11;
+namespace Aidge {
+
+void init_TanhClamp(py::module& m) {
+    py::class_<TanhClamp_Op, std::shared_ptr<TanhClamp_Op>, OperatorTensor>(m, "TanhClampOp", py::multiple_inheritance())
+    .def(py::init<>())
+    .def_static("get_inputs_name", &TanhClamp_Op::getInputsName)
+    .def_static("get_outputs_name", &TanhClamp_Op::getOutputsName);
+    declare_registrable<TanhClamp_Op>(m, "TanhClampOp");
+    m.def("TanhClamp", &TanhClamp, py::arg("name") = "");
+}
+}  // namespace Aidge

python_binding/operator/pybind_FixedQ.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/FixedQ.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+
+namespace py = pybind11;
+namespace Aidge {
+
+void init_FixedQ(py::module& m) {
+    py::class_<FixedQ_Op, std::shared_ptr<FixedQ_Op>, OperatorTensor>(m, "FixedQOp", py::multiple_inheritance())
+    .def(py::init<std::size_t, float, bool>(), py::arg("nb_bits"), py::arg("span"), py::arg("is_output_unsigned"))
+    .def_static("get_inputs_name", &FixedQ_Op::getInputsName)
+    .def_static("get_outputs_name", &FixedQ_Op::getOutputsName);
+    declare_registrable<FixedQ_Op>(m, "FixedQOp");
+    m.def("FixedQ", &FixedQ, py::arg("nb_bits") = 8, py::arg("span") = 4.0f, py::arg("is_output_unsigned") = false, py::arg("name") = "");
+}
+}  // namespace Aidge
\ No newline at end of file

python_binding/operator/pybind_LSQ.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/operator/LSQ.hpp"
+#include "aidge/operator/OperatorTensor.hpp"
+
+namespace py = pybind11;
+namespace Aidge {
+
+void init_LSQ(py::module& m) {
+    py::class_<LSQ_Op, std::shared_ptr<LSQ_Op>, OperatorTensor>(m, "LSQOp", py::multiple_inheritance())
+    .def(py::init<const std::pair<int, int>&>(), py::arg("range") = std::pair<int, int>{0, 255})
+    .def_static("get_inputs_name", &LSQ_Op::getInputsName)
+    .def_static("get_outputs_name", &LSQ_Op::getOutputsName);
+    declare_registrable<LSQ_Op>(m, "LSQOp");
+    m.def("LSQ", &LSQ, py::arg("range") = std::pair<int, int>{0, 255}, py::arg("name") = "");
+}
+}  // namespace Aidge

python_binding/pybind_PTQ.cpp

@@ -14,9 +14,9 @@
 
 #include <string>
 
-#include "aidge/PTQ/Clip.hpp"
-#include "aidge/PTQ/CLE.hpp"
-#include "aidge/PTQ/PTQ.hpp"
+#include "aidge/quantization/PTQ/Clip.hpp"
+#include "aidge/quantization/PTQ/CLE.hpp"
+#include "aidge/quantization/PTQ/PTQ.hpp"
 
 #include "aidge/hook/Hook.hpp"
 #include "aidge/graph/GraphView.hpp"
@@ -24,7 +24,7 @@
 namespace py = pybind11;
 
 namespace Aidge {
-void init_QuantPTQ(py::module &m) {
+void init_PTQ(py::module &m) {
 
     py::enum_<Clipping>(m, "Clipping", "Kind of clipping policy to apply during the activation quantization")
         .value("MAX", Clipping::MAX)
@@ -32,7 +32,7 @@ void init_QuantPTQ(py::module &m) {
         .value("AA" , Clipping::AA)
         .value("KL" , Clipping::KL);
 
-    m.def("check_architecture", &checkArchitecture, py::arg("network"),     
+    m.def("check_architecture", &checkArchitecture, py::arg("network"),
     R"mydelimiter(
     Determine whether an input GraphView can be quantized or not.
     :param network: The GraphView to be checked.
@@ -41,15 +41,15 @@ void init_QuantPTQ(py::module &m) {
     :rtype: bool
     )mydelimiter");
 
-    m.def("insert_scaling_nodes", &insertScalingNodes, py::arg("network"),     
+    m.def("insert_scaling_nodes", &insertScalingNodes, py::arg("network"),
     R"mydelimiter(
-    Insert a scaling node after each affine node of the GraphView. 
+    Insert a scaling node after each affine node of the GraphView.
     Also insert a scaling node in every purely residual branches.
     :param network: The GraphView containing the affine nodes.
     :type network: :py:class:`aidge_core.GraphView`
     )mydelimiter");
 
-    m.def("normalize_parameters", &normalizeParameters, py::arg("network"),    
+    m.def("normalize_parameters", &normalizeParameters, py::arg("network"),
     R"mydelimiter(
     Normalize the parameters of each parametrized node, so that they fit in the [-1:1] range.
     :param network: The GraphView containing the parametrized nodes.
@@ -70,17 +70,17 @@ void init_QuantPTQ(py::module &m) {
     )mydelimiter");
 
     m.def("normalize_activations", &normalizeActivations, py::arg("network"), py::arg("value_ranges"),
-    R"mydelimiter(      
+    R"mydelimiter(
     Normalize the activations of each affine node so that they fit in the [-1:1] range.
     This is done by reconfiguring the scaling nodes, as well as rescaling the weights and biases tensors.
     :param network: The GraphView containing the affine nodes.
     :type network: :py:class:`aidge_core.GraphView`
-    :param value_ranges: The node output value ranges computed over the calibration dataset. 
+    :param value_ranges: The node output value ranges computed over the calibration dataset.
     :type value_ranges: list of float.
     )mydelimiter");
 
-    m.def("quantize_normalized_network", &quantizeNormalizedNetwork, py::arg("network"), py::arg("nb_bits"), py::arg("apply_rounding"), py::arg("optimize_signs"), py::arg("verbose"),
-    R"mydelimiter( 
+    m.def("quantize_normalized_network", &quantizeNormalizedNetwork, py::arg("network"), py::arg("nb_bits"), py::arg("apply_rounding"), py::arg("optimize_signs"), py::arg("verbose") = false,
+    R"mydelimiter(
     Quantize an already normalized (in term of parameters and activations) network.
     :param network: The GraphView to be quantized.
     :type network: :py:class:`aidge_core.GraphView`
@@ -94,14 +94,14 @@ void init_QuantPTQ(py::module &m) {
     :type verbose: bool
     )mydelimiter");
 
-    m.def("quantize_network", &quantizeNetwork ,py::arg("network"), py::arg("nb_bits"), py::arg("input_dataset"), py::arg("clipping_mode") = "MAX", py::arg("apply_rounding") = true, py::arg("optimize_signs") = false, py::arg("single_shift") = false, py::arg("verbose") = false,
-    R"mydelimiter(   
+    m.def("quantize_network", &quantizeNetwork ,py::arg("network"), py::arg("nb_bits"), py::arg("input_dataset"), py::arg("clipping_mode") = Clipping::MAX, py::arg("apply_rounding") = true, py::arg("optimize_signs") = false, py::arg("single_shift") = false, py::arg("verbose") = false,
+    R"mydelimiter(
     Main quantization routine. Performs every step of the quantization pipeline.
     :param network: The GraphView to be quantized.
     :type network: :py:class:`aidge_core.GraphView`
     :param nb_bits: The desired number of bits of the quantization.
     :type nb_bits: int
-    :param input_dataset: The input dataset on which the value ranges are computed. 
+    :param input_dataset: The input dataset on which the value ranges are computed.
     :type input_dataset: list of :py:class:`aidge_core.Tensor`
     :param clipping_mode: Type of the clipping optimization. Can be either 'MAX', 'MSE', 'AA' or 'KL'.
     :type clipping_mode: string
@@ -109,13 +109,13 @@ void init_QuantPTQ(py::module &m) {
     :type apply_rounding: bool
     :param optimize_signs: Whether to take account of the IO signs of the operators or not.
     :type optimize_signs: bool
-    :param single_shift: Whether to convert the scaling factors into powers of two. If true the approximations are compensated using the previous nodes weights. 
+    :param single_shift: Whether to convert the scaling factors into powers of two. If true the approximations are compensated using the previous nodes weights.
     :type single_shift: bool
     :param verbose: Whether to print internal informations about the quantization process.
     :type verbose: bool
     )mydelimiter");
 
-    m.def("compute_histograms", &computeHistograms, py::arg("value_ranges"), py::arg("nb_bins"), py::arg("network"), py::arg("input_dataset"), 
+    m.def("compute_histograms", &computeHistograms, py::arg("value_ranges"), py::arg("nb_bins"), py::arg("network"), py::arg("input_dataset"),
     R"mydelimiter(
     Compute the histograms of the activations of each node contained in the map of the ranges (passed as argument).
     :param value_ranges: A map associating each considered node name to its corresponding output range.
@@ -154,11 +154,11 @@ void init_QuantPTQ(py::module &m) {
     :rtype: float
     )mydelimiter");
 
-    m.def("adjust_ranges", &adjustRanges, py::arg("clipping_mode"), py::arg("value_ranges"), py::arg("nb_bits"), py::arg("network"), py::arg("input_dataset"), py::arg("verbose"),
+    m.def("adjust_ranges", &adjustRanges, py::arg("clipping_mode"), py::arg("value_ranges"), py::arg("nb_bits"), py::arg("network"), py::arg("input_dataset"), py::arg("verbose") = false,
     R"mydelimiter(
-    Return a corrected map of the provided activation ranges. 
-    To do so compute the optimal clipping values for every node and multiply the input ranges by those values. 
-    The method used to compute the clippings can be eihter 'MSE', 'AA', 'KL' or 'MAX'. 
+    Return a corrected map of the provided activation ranges.
+    To do so compute the optimal clipping values for every node and multiply the input ranges by those values.
+    The method used to compute the clippings can be eihter 'MSE', 'AA', 'KL' or 'MAX'.
     :param clipping_mode: The method used to compute the optimal clippings.
     :type clipping_mode: enum
     :param value_ranges: The map associating each affine node to its output range.
@@ -176,30 +176,30 @@ void init_QuantPTQ(py::module &m) {
     )mydelimiter");
 
 
-    m.def("compute_sign_map", &computeSignMap, py::arg("network"), py::arg("verbose"),
+    m.def("compute_sign_map", &computeSignMap, py::arg("network"), py::arg("verbose") = false,
     R"mydelimiter(
-    For each node, compute the sign of its input and output values. 
+    For each node, compute the sign of its input and output values.
     The goal of the routine is to maximize the number of unsigned IOs in order to double the value resolution when possible.
     :param network: The GraphView to analyze.
-    :type network: :py:class:`aidge_core.GraphView`  
+    :type network: :py:class:`aidge_core.GraphView`
     :param verbose: Whether to print the sign map or not.
-    :type verbose: bool  
+    :type verbose: bool
     :return: A map associating a pair of signs to each node of the GraphView (a sign for the input and one for the output).
-    :rtype: dict  
+    :rtype: dict
     )mydelimiter");
 
     m.def("cross_layer_equalization", &crossLayerEqualization, py::arg("network"), py::arg("target_delta"),
     R"mydelimiter(
-    Equalize the ranges of the nodes parameters by proceding iteratively. 
+    Equalize the ranges of the nodes parameters by proceding iteratively.
     Can only be applied to single branch networks (otherwise does not edit the graphView).
     :param network: The GraphView to process.
-    :type network: :py:class:`aidge_core.GraphView`  
+    :type network: :py:class:`aidge_core.GraphView`
     :param target_delta: the stopping criterion (typical value : 0.01)
-    :type target_delta: float  
+    :type target_delta: float
     )mydelimiter");
 
     m.def("get_weight_ranges", &getWeightRanges, py::arg("network"),
-    R"mydelimiter(  
+    R"mydelimiter(
     Compute the weight ranges of every affine nodes. Provided for debugging purposes.
     :param network: The GraphView containing the affine nodes.
     :type network: :py:class:`aidge_core.GraphView`
@@ -207,23 +207,22 @@ void init_QuantPTQ(py::module &m) {
     :rtype: dict
     )mydelimiter");
 
-    m.def("clear_biases", &clearBiases, py::arg("network"),   
+    m.def("clear_biases", &clearBiases, py::arg("network"),
     R"mydelimiter(
     Clear the affine nodes biases. Provided form debugging purposes.
     :param network: The GraphView containing the affine nodes.
     :type network: :py:class:`aidge_core.GraphView`
     )mydelimiter");
 
-    m.def("dev_ptq", &devPTQ, py::arg("network"), 
+    m.def("dev_ptq", &devPTQ, py::arg("network"),
     R"mydelimiter(
     Developement and test routine.
     :param network: The GraphView under test.
-    :type network: :py:class:`aidge_core.GraphView`    
+    :type network: :py:class:`aidge_core.GraphView`
     )mydelimiter");
-}
 
-PYBIND11_MODULE(aidge_quantization, m) {
-    init_QuantPTQ(m);
+    m.def("prepare_network", &prepareNetwork, py::arg("network"), "prepare the network fo the PTQ");
+
 }
 
 } // namespace Aidge

python_binding/pybind_QAT_FixedQ.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "aidge/quantization/QAT/QAT_FixedQ.hpp"
+#include "aidge/hook/Hook.hpp"
+#include "aidge/graph/GraphView.hpp"
+
+namespace py = pybind11;
+
+namespace Aidge {
+
+void init_QAT_FixedQ(py::module &m) {
+
+    auto mQuantFixedQ = m.def_submodule("fixedq");
+
+    mQuantFixedQ.def("insert_quantizers", &QuantFixedQ::insertQuantizers, py::arg("network"), py::arg("nb_bits"), py::arg("span"));
+
+    mQuantFixedQ.def("insert_and_init_quantizers", &QuantFixedQ::insertAndInitQuantizers, py::arg("network"), py::arg("nb_bits"), py::arg("calibration_data"), py::arg("scale"));
+
+    mQuantFixedQ.def("dev_qat", &QuantFixedQ::devQAT, py::arg("network"));
+}
+} // namespace Aidge

python_binding/pybind_QAT_LSQ.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "aidge/quantization/QAT/QAT_LSQ.hpp"
+#include "aidge/hook/Hook.hpp"
+#include "aidge/graph/GraphView.hpp"
+
+namespace py = pybind11;
+
+namespace Aidge {
+
+void init_QAT_LSQ(py::module &m) {
+
+    auto mQuantLSQ = m.def_submodule("lsq");
+
+    mQuantLSQ.def("insert_quantizers", &QuantLSQ::insertQuantizers, py::arg("network"), py::arg("nb_bits"), py::arg("step_size"));
+
+    mQuantLSQ.def("insert_and_init_quantizers", &QuantLSQ::insertAndInitQuantizers, py::arg("network"), py::arg("nb_bits"), py::arg("calibration_data"));
+}
+} // namespace Aidge

python_binding/pybind_Quantization.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include "aidge/backend/QuantizationCPU.hpp"
+//#include "aidge/backend/QuantizationCUDA.hpp"
+
+namespace py = pybind11;
+
+namespace Aidge 
+{
+
+// operators
+void init_FixedQ(py::module& m);
+void init_LSQ(py::module& m);
+void init_TanhClamp(py::module& m);
+void init_DoReFa(py::module& m);
+
+
+// quantization routines
+void init_PTQ(py::module &m);
+void init_QAT_FixedQ(py::module &m);
+void init_QAT_LSQ(py::module &m);
+
+
+PYBIND11_MODULE(aidge_quantization, m) 
+{
+    init_FixedQ(m);
+    init_LSQ(m);
+    init_TanhClamp(m);
+    init_DoReFa(m);
+
+    init_PTQ(m);
+    init_QAT_FixedQ(m);
+    init_QAT_LSQ(m);
+}
+
+} // namespace Aidge

setup.py

 #!/usr/bin/env python3
-""" Aidge
-
-#TODO To change
-POC of the next framework named Aidge
-"""
-
-DOCLINES = (__doc__ or '').split("\n")
-
 import sys
 import os
 
-# Python supported version checks
-if sys.version_info[:2] < (3, 7):
-    raise RuntimeError("Python version >= 3.7 required.")
-
-
-CLASSIFIERS = """\
-Development Status :: 2 - Pre-Alpha
-"""
-
 import shutil
 import pathlib
-import subprocess
 import multiprocessing
 
 from math import ceil
 
+import toml
+
 from setuptools import setup, Extension
-from setuptools import find_packages
 from setuptools.command.build_ext import build_ext
 
+
 def get_project_name() -> str:
-    return open(pathlib.Path().absolute() / "project_name.txt", "r").read()
+    with open(pathlib.Path().absolute() / "pyproject.toml", "r") as file:
+        project_toml = toml.load(file)
+        return project_toml["project"]["name"]
+
 
 def get_project_version() -> str:
     aidge_root = pathlib.Path().absolute()
@@ -43,8 +30,8 @@ class CMakeExtension(Extension):
     def __init__(self, name):
         super().__init__(name, sources=[])
 
-class CMakeBuild(build_ext):
 
+class CMakeBuild(build_ext):
     def run(self):
         # This lists the number of processors available on the machine
         # The compilation will use half of them
@@ -62,17 +49,48 @@ class CMakeBuild(build_ext):
 
         os.chdir(str(build_temp))
 
-        # Impose to use the executable of the python
-        # used to launch setup.py to setup PythonInterp
-        param_py = "-DPYTHON_EXECUTABLE=" + sys.executable
-
-        install_path = os.path.join(sys.prefix, "lib", "libAidge")  if "AIDGE_INSTALL" not in os.environ else os.environ["AIDGE_INSTALL"]
+        python_executable = sys.executable
+        print(f"python executable : {python_executable}")
+
+        compile_type = (
+            "Release"
+            if "AIDGE_PYTHON_BUILD_TYPE" not in os.environ
+            else os.environ["AIDGE_PYTHON_BUILD_TYPE"]
+        )
+
+        install_path = (
+            os.path.join(sys.prefix, "lib", "libAidge")
+            if "AIDGE_INSTALL" not in os.environ
+            else os.environ["AIDGE_INSTALL"]
+        )
+
+        build_gen = (
+            ["-G", os.environ["AIDGE_BUILD_GEN"]]
+            if "AIDGE_BUILD_GEN" in os.environ
+            else []
+        )
+
+        self.spawn(
+            [
+                "cmake",
+                *build_gen,
+                str(cwd),
+                "-DTEST=OFF",
+                f"-DCMAKE_INSTALL_PREFIX:PATH={install_path}",
+                f"-DCMAKE_BUILD_TYPE={compile_type}",
+                "-DPYBIND=ON",
+                "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
+                "-DCOVERAGE=OFF",
+            ]
+        )
 
-        self.spawn(['cmake', str(cwd), param_py, '-DTEST=OFF', f'-DCMAKE_INSTALL_PREFIX:PATH={install_path}'])
         if not self.dry_run:
-            self.spawn(['make', 'all', 'install', '-j', max_jobs])
+            self.spawn(
+                ["cmake", "--build", ".", "--config", compile_type, "-j", max_jobs]
+            )
+            self.spawn(["cmake", "--install", ".", "--config", compile_type])
         os.chdir(str(cwd))
-
+        
         aidge_package = build_lib / (get_project_name())
 
         # Get "aidge core" package
@@ -81,8 +99,8 @@ class CMakeBuild(build_ext):
         # Copy all shared object files from build_temp/lib to aidge_package
         for root, _, files in os.walk(build_temp.absolute()):
             for file in files:
-                if file.endswith('.so') and (root != str(aidge_package.absolute())):
-                    currentFile=os.path.join(root, file)
+                if file.endswith(".so") and (root != str(aidge_package.absolute())):
+                    currentFile = os.path.join(root, file)
                     shutil.copy(currentFile, str(aidge_package.absolute()))
 
         # Copy version.txt in aidge_package
@@ -90,23 +108,12 @@ class CMakeBuild(build_ext):
         shutil.copy("version.txt", str(aidge_package.absolute()))
 
 
-if __name__ == '__main__':
-
+if __name__ == "__main__":
     setup(
-        name=get_project_name(),
-        version=get_project_version(),
-        python_requires='>=3.7',
-        description=DOCLINES[0],
-        long_description_content_type="text/markdown",
-        long_description="\n".join(DOCLINES[2:]),
-        classifiers=[c for c in CLASSIFIERS.split('\n') if c],
-        platforms=["Linux"],
-        packages=find_packages(where="."),
         include_package_data=True,
         ext_modules=[CMakeExtension(get_project_name())],
         cmdclass={
-            'build_ext': CMakeBuild,
+            "build_ext": CMakeBuild,
         },
         zip_safe=False,
-
     )

src/PTQ/CLE.cpp

@@ -84,11 +84,11 @@ void crossLayerEqualization(std::shared_ptr<GraphView> graphView, float targetDe
     do 
     {
         maxRangeDelta = 0.0;
-        /*
-        std::cout << " ----- " << std::endl;
-        for (std::shared_ptr<Node> node : affineNodeVector)
-            std::cout << getTensorAbsoluteMax(getWeightTensor(node)) << std::endl;
-        */
+        
+        //std::cout << " ----- " << std::endl;
+        //for (std::shared_ptr<Node> node : affineNodeVector)
+        //    std::cout << getTensorAbsoluteMax(getWeightTensor(node)) << std::endl;
+        
         for (size_t i = 0; i < (affineNodeVector.size() - 1); i++)
         {
             std::shared_ptr<Node> n1 = affineNodeVector[i];

src/PTQ/Clip.cpp

@@ -53,7 +53,7 @@ std::map<std::string, std::vector<int>> computeHistograms(std::map<std::string,
 
     for (std::shared_ptr<Tensor> inputTensor : inputDataSet)
     {
-        Log::info(" IT (BIS) : {}", it++);
+        Log::debug(" IT (BIS) : {}", it++);
 
         // Inference ...
 
@@ -225,4 +225,4 @@ std::map<std::string, float> adjustRanges(Clipping clippingMode, std::map<std::s
     return valueRanges;
 }
 
-}
\ No newline at end of file
+}

src/PTQ/PTQ.cpp

@@ -69,7 +69,7 @@ bool isMerging(std::shared_ptr<Node> node)
 
 bool checkArchitecture(std::shared_ptr<GraphView> graphView)
 {
-    std::set<std::string> otherNodeTypes({"Flatten", "Softmax", "ReLU", "Producer"});
+    std::set<std::string> otherNodeTypes({"Flatten", "Softmax", "BatchNorm2D", "ReLU", "Producer"});
 
     for (std::shared_ptr<Node> node : graphView->getNodes())
     {
@@ -128,8 +128,7 @@ static float getTensorAbsoluteMax(std::shared_ptr <Tensor> tensor)
     return maxValue;
 }
 
-static void removeMatchingNodes(std::vector<std::shared_ptr<Node>>& nodeVector, std::string nodeType)
-{
+/*
     std::vector<std::shared_ptr<Node>>::iterator iter = nodeVector.begin();
     while (iter != nodeVector.end())
     {
@@ -138,6 +137,17 @@ static void removeMatchingNodes(std::vector<std::shared_ptr<Node>>& nodeVector,
         else
            ++iter;
     }
+*/
+
+// TODO : pass nodeVector by reference ...
+static std::vector<std::shared_ptr<Node>> removeMatchingNodes(std::vector<std::shared_ptr<Node>> nodeVector, std::string nodeType)
+{
+    std::vector<std::shared_ptr<Node>> remainingNodes;
+    for (std::shared_ptr<Node> node : nodeVector)
+        if (node->type() != nodeType)
+            remainingNodes.push_back(node);
+
+    return remainingNodes;
 }
 
 static void fixScheduling(std::vector<std::shared_ptr<Node>>& nodeVector) {
@@ -181,10 +191,14 @@ std::vector<std::shared_ptr<Node>> retrieveNodeVector(std::shared_ptr<GraphView>
 
     nodeVector = scheduler.getStaticScheduling();
 
+    //std::cout << " RNV : NB OF NODES = " << nodeVector.size() << std::endl;
+    //for (auto node : nodeVector)
+    //    std::cout << node->type() << std::endl;
+
     fixScheduling(nodeVector);
-    removeMatchingNodes(nodeVector, "Producer");
+    nodeVector = removeMatchingNodes(nodeVector, "Producer");
 
-    if (verbose)
+    if (verbose) 
     {
         Log::info("NB OF NODES = {}", nodeVector.size());
         for (std::shared_ptr<Node> node : nodeVector)
@@ -216,7 +230,7 @@ static void popSoftMax(std::shared_ptr<GraphView> graphView)
     }
 }
 
-static void prepareNetwork(std::shared_ptr<GraphView> graphView)
+void prepareNetwork(std::shared_ptr<GraphView> graphView)
 {
     removeFlatten(graphView);
 
@@ -339,6 +353,18 @@ static std::shared_ptr<Node> getPreviousScalingNode(std::shared_ptr<Node> mergin
     return currNode;
 }
 
+// XXX double check this !
+static bool nodeHasBias(std::shared_ptr<Node> node)
+{
+    if (node->getParents().size() == 3)
+    {
+        std::shared_ptr<Tensor> biasTensor = getBiasTensor(node);
+        if (biasTensor)
+            return true;
+    }
+    return false;
+}
+
 void normalizeParameters(std::shared_ptr<GraphView> graphView)
 {
     // CREATE THE ACCUMULATED RATIO MAP ///////////////////////////////////////
@@ -387,15 +413,12 @@ void normalizeParameters(std::shared_ptr<GraphView> graphView)
                 accumulatedRatios[node->name()] = accumulatedRatios[prevNode->name()] * ratio;
             }
 
-            // Handle the bias ...
-            bool nodeHasBias = (node->getParents().size() == 3);
-            if (nodeHasBias)
+            // Handle the bias ..
+
+            if (nodeHasBias(node))
             {
                 std::shared_ptr<Tensor> biasTensor = getBiasTensor(node);
-
-                // Check that a bias is present (as it is optional)
-                if (biasTensor)
-                    rescaleTensor(biasTensor, accumulatedRatios[node->name()] );
+                rescaleTensor(biasTensor, accumulatedRatios[node->name()] );
             }
         }
 
@@ -478,7 +501,7 @@ std::map<std::string, float> computeRanges(std::shared_ptr<GraphView> graphView,
 
     for (std::shared_ptr<Tensor> sample : inputDataSet)
     {
-        Log::info(" IT : {}", it++);
+        Log::debug(" IT : {}", it++);
 
         // Inference ...
 
@@ -569,8 +592,9 @@ void normalizeActivations(std::shared_ptr<GraphView> graphView, std::map<std::st
 
             if (isAffine(prevNode))
             {
-                bool prevNodeHasBias = (prevNode->getParents().size() == 3);
-                if (prevNodeHasBias)  {
+                bool prevNodeHasBias = nodeHasBias(prevNode);
+                if (prevNodeHasBias)  
+                {
                     std::shared_ptr<Tensor> biasTensor = getBiasTensor(prevNode);
                     rescaleTensor(biasTensor, 1.0 / prevScalingFactor);
                 }
@@ -722,14 +746,14 @@ std::map<std::string, std::pair<bool, bool>> computeSignMap(std::shared_ptr<Grap
     // SANITY CHECK (TEMPORARY)
 
     for (std::shared_ptr<Node> node : nodeVector)
-        if (node != firstNode)
+    {
+        for (std::shared_ptr<Node> child : node->getChildren()) 
         {
-            for (std::shared_ptr<Node> parent : node->getParents()) 
-                if (parent->type() != "Producer")
-                    if (signMap[parent->name()].second != signMap[node->name()].first)
-                        Log::error(" computeSignMap : link is not sane ! ({} -> {})", parent->name(), node->name());
+            if (signMap[node->name()].second != signMap[child->name()].first)
+                Log::error(" computeSignMap : link is not sane ! ({} -> {})", node->name(), child->name());                    
         }
-
+    }
+    
     return signMap;
 }
 
@@ -769,8 +793,7 @@ void quantizeNormalizedNetwork(std::shared_ptr<GraphView> graphView, std::uint8_
 
             // Rescale the bias tensor
 
-            bool nodeHasBias = (node->getParents().size() == 3);
-            if (nodeHasBias)  
+            if (nodeHasBias(node))  
             {
                 bool inputIsUnsigned = signMap[node->name()].first;
                 float rescaling = inputIsUnsigned ? unsignedMax * signedMax : signedMax * signedMax;
@@ -919,8 +942,7 @@ void performSingleShiftApproximation(std::shared_ptr<GraphView> graphView, bool
             if (applyRounding)
                 roundTensor(weightTensor);
 
-            bool nodeHasBias = (node->getParents().size() == 3);
-            if (nodeHasBias)
+            if (nodeHasBias(node))
             {
                 std::shared_ptr<Tensor> biasTensor = getBiasTensor(node);
                 rescaleTensor(biasTensor, ratio);

src/QAT/QAT_FixedQ.cpp

+/********************************************************************************
+ * Copyright (c) 2023 CEA-List
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-2.0.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ ********************************************************************************/
+
+#include "aidge/quantization/QAT/QAT_FixedQ.hpp"
+#include "aidge/operator/FixedQ.hpp"
+
+#include "aidge/data/Tensor.hpp"
+#include "aidge/graph/GraphView.hpp"
+#include "aidge/scheduler/SequentialScheduler.hpp"
+#include "aidge/scheduler/Scheduler.hpp"
+#include "aidge/graph/Matching.hpp"
+
+namespace Aidge {
+
+void QuantFixedQ::insertQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, float span)
+{
+    const auto matches = SinglePassGraphMatching(graphView).match("(Conv#|FC#)");
+
+    for (const auto& match : matches) 
+    {
+        auto linearNode = match.graph->rootNode(); 
+
+        // INPUT QUANTIZERS INSERTION
+
+        auto inputQuantizerName = linearNode->name() + "_fixedq_i";  // TODO : double check this, and use createUniqueName()
+        auto inputQuantizerNode = FixedQ(nbBits, span, false, inputQuantizerName);
+
+        // Absorb the ReLU when possible ...
+
+        bool nodeHasParent = static_cast<bool> (linearNode->getParents()[0]); // XXX is this safe ???
+
+        if (nodeHasParent) {
+            auto parentNode = linearNode->getParents()[0];
+            if (parentNode->type() == "ReLU") {
+                auto inputQuantizerOp = std::static_pointer_cast<FixedQ_Op> (inputQuantizerNode->getOperator());
+                inputQuantizerOp->isOutputUnsigned() = true;
+                graphView->replace({parentNode}, {}); 
+            }
+        }
+
+        // We need to handle the case where the linear node is the first one ...
+
+        if (nodeHasParent) {
+            graphView->insertParent(linearNode, inputQuantizerNode, 0, 0, 0);
+        } else {
+            inputQuantizerNode->addChild(graphView);
+            graphView->add(inputQuantizerNode);
+        }
+
+        // PARAM QUANTIZERS INSERTION
+
+        auto paramQuantizerName = linearNode->name() + "_fixedq_p";  // TODO : double check this, and use createUniqueName()
+        auto paramQuantizerNode = FixedQ(nbBits, span, false, paramQuantizerName); 
+        graphView->insertParent(linearNode, paramQuantizerNode, 1, 0, 0);
+    }
+}
+
+static float getTensorStd(std::shared_ptr<Tensor> tensor)
+{
+    float acc = 0;
+    float * castedTensor = static_cast<float *> (tensor->getImpl()->rawPtr());
+    for(std::size_t i = 0; i < tensor->size(); i++)
+        acc += castedTensor[i] * castedTensor[i];
+    acc /= static_cast<float> (tensor->size());
+    return std::sqrt(acc);
+}
+
+static std::map<std::string, float> collectInputStats(std::shared_ptr<GraphView> graphView, std::shared_ptr<Tensor> calibrationData)
+{
+    // Propagate the calibration tensor
+
+    SequentialScheduler scheduler(graphView);
+    scheduler.resetScheduling();
+    scheduler.forward(true, {calibrationData});
+
+    // Store the input tensor statistics
+
+    std::map<std::string, float> inputStats;
+    for (auto node : graphView->getNodes())
+    {
+        if (node->type() == "FC" || node->type() == "Conv") // TODO: use graph matching !!!
+        {
+            const auto op = std::static_pointer_cast<FixedQ_Op>(node->getOperator());
+            float inputStd = getTensorStd(op->getInput(0));
+            inputStats.insert(std::make_pair(node->name(), inputStd));
+            std::cout << node->name() << " -> " << inputStd << std::endl;
+        }
+    }
+
+    return inputStats;
+}
+
+static std::map<std::string, float> collectParamStats(std::shared_ptr<GraphView> graphView)
+{
+    std::map<std::string, float> paramStats;
+    for (auto node : graphView->getNodes())
+    {
+        if (node->type() == "FC" || node->type() == "Conv") // TODO: use graph matching !!!
+        {
+            const auto op = std::static_pointer_cast<FixedQ_Op>(node->getOperator());
+            float paramStd = getTensorStd(op->getInput(1));
+            paramStats.insert(std::make_pair(node->name(), paramStd));
+            std::cout << node->name() << " -> " << paramStd << std::endl;
+        }
+    }
+    
+    return paramStats;
+}
+
+static void adjustQuantizersSpans(std::shared_ptr<GraphView> graphView, std::map<std::string, float> inputStats, std::map<std::string, float> paramStats, float scale = 4.0f)
+{
+    const auto matches = SinglePassGraphMatching(graphView).match("(Conv#|FC#)");
+
+    for (const auto& match : matches) 
+    {
+        auto linearNode = match.graph->rootNode();
+
+        // Adjust the input quantizers spans
+
+        auto inputQuantNode = linearNode->getParent(0);
+        auto inputQuantOp = std::static_pointer_cast<FixedQ_Op>(inputQuantNode->getOperator());
+        inputQuantOp->span() = inputStats[linearNode->name()] * scale;
+
+        // Adjust the param quantizers spans
+
+        auto paramQuantNode = linearNode->getParent(1);
+        auto paramQuantOp = std::static_pointer_cast<FixedQ_Op>(paramQuantNode->getOperator());
+        paramQuantOp->span() = paramStats[linearNode->name()] * scale;
+    }
+}
+
+void QuantFixedQ::insertAndInitQuantizers(std::shared_ptr<GraphView> graphView, size_t nbBits, std::shared_ptr<Tensor> calibrationData, float scale)
+{
+    // Collect the tensor statisics
+    auto inputStats = collectInputStats(graphView, calibrationData);
+    auto paramStats = collectParamStats(graphView);
+
+    // Insert the quantizers
+    insertQuantizers(graphView, nbBits, 1.0);
+
+    // Adjust the quantizers spans
+    adjustQuantizersSpans(graphView, inputStats, paramStats, scale);
+}
+
+void QuantFixedQ::devQAT(std::shared_ptr<GraphView> graphView) 
+{
+    SequentialScheduler scheduler(graphView);
+    scheduler.generateScheduling();
+    auto s = scheduler.getStaticScheduling();
+    for (std::shared_ptr<Node> node : s)
+        std::cout << " name : " << node->name() << std::endl;
+}
+
+}
\ No newline at end of file