Skip to content
Snippets Groups Projects
Commit 7c3813bb authored by Gallas Gaye's avatar Gallas Gaye
Browse files

feat: Softmax works with any number of dimensions

parent 203fee0d
Branches issue#199
No related tags found
2 merge requests!39Update 0.2.1 -> 0.3.0,!36feat: Add missing operators for AIDGE model benchmarking
......@@ -6,50 +6,48 @@
#include "kernels/macs.hpp"
#include <type_traits>
#include <cmath>
#include <algorithm>
template<int NB_CHANNELS,
int CHANNELS_HEIGHT, int CHANNELS_WIDTH,
int NB_OUTPUTS,
int OUTPUTS_HEIGHT, int OUTPUTS_WIDTH,
int AXIS,
template<int AXIS_SIZE,
int AXIS_SIZE_POST,
int AXIS_SIZE_PRE,
typename Input_T, typename Output_T>
__attribute__((always_inline)) inline
void softmax_forward (
const Input_T* __restrict inputs,
Output_T* __restrict outputs)
{
Input_T maxValue = 0.0f;
for (int och = 0; och < NB_OUTPUTS; och++) {
maxValue = std::max(maxValue, inputs[och]);
}
Input_T sumExp = 0.0f;
if constexpr (std::is_same_v<Input_T, Output_T>) {
for (int och = 0; och < NB_OUTPUTS; och++) {
// This should be both more performant while keeping the same memory footprint but we can only use it if INPUT_T and OUTPUT_T types are the same !
outputs[och] = std::exp(inputs[och] - maxValue);
sumExp += outputs[och];
}
for (int och = 0; och < NB_OUTPUTS; och++) {
outputs[och] /= sumExp;
}
}
else
{
for (int och = 0; och < NB_OUTPUTS; och++) {
sumExp += std::exp(inputs[och] - maxValue);
}
for (int och = 0; och < NB_OUTPUTS; och++) {
outputs[och] = std::exp(inputs[och] - maxValue) / sumExp;
// Iterate over the "pre-axis" and "post-axis" slices.
// For each slice along the axis, compute the maximum value,
// the sum of exponentials, and then write the normalized softmax outputs.
for (int i = 0; i < AXIS_SIZE_PRE; ++i) {
for (int j = 0; j < AXIS_SIZE_POST; ++j) {
// Compute the base index for this slice.
const int baseIdx = i * AXIS_SIZE * AXIS_SIZE_POST + j;
// Find the maximum value along the axis.
Input_T maxVal = inputs[baseIdx];
for (int k = 1; k < AXIS_SIZE; ++k) {
const int idx = baseIdx + k * AXIS_SIZE_POST;
maxVal = std::max(maxVal, inputs[idx]);
}
// Compute the sum of the exponentials along the axis.
Input_T sumExp = 0;
for (int k = 0; k < AXIS_SIZE; ++k) {
const int idx = baseIdx + k * AXIS_SIZE_POST;
outputs[idx] = std::exp(inputs[idx] - maxVal);
sumExp += outputs[idx];
}
// Write the softmax values to the output.
for (int k = 0; k < AXIS_SIZE; ++k) {
const int idx = baseIdx + k * AXIS_SIZE_POST;
outputs[idx] /= sumExp;
}
}
}
}
#endif // __AIDGE_EXPORT_CPP_KERNELS_SOFTMAX__
......@@ -338,6 +338,30 @@ class SoftmaxCPP(ExportNodeCpp):
def __init__(self, node, mem_info):
super().__init__(node, mem_info)
self.attributes["axis"] = node.get_operator().attr.axis
assert self.node.get_nb_inputs() == 1, (
f"export softmax: nb_inputs == {self.node.get_nb_inputs()} not implemented"
)
tensor = self.operator.get_input(0)
nbDims = len(tensor.dims())
assert self.attributes["axis"] < nbDims, (
f"export softmax: attribute axis == {node.get_operator().attr.axis} should be less than {nbDims}"
)
postAxisElems = 1
for i in range(self.attributes["axis"] + 1, nbDims):
postAxisElems *= tensor.dims()[i]
preAxisElems = 1
for i in range(self.attributes["axis"]):
preAxisElems *= tensor.dims()[i]
self.attributes["axis_size"] = tensor.dims()[self.attributes["axis"]]
self.attributes["axis_size_post"] = postAxisElems
self.attributes["axis_size_pre"] = preAxisElems
self.config_template = str(
ROOT / "templates" / "configuration" / "softmax_config.jinja")
self.forward_template = str(
......
......@@ -7,6 +7,8 @@
{#- Calculate sizes #}
{%- set weights_size = out_chan[0] * in_chan[0] * in_height[0] * in_width[0] %}
#define {{ name|upper }}_AXIS {{ axis }}
#define {{ name|upper }}_AXIS_SIZE {{ axis_size }}
#define {{ name|upper }}_AXIS_SIZE_POST {{ axis_size_post }}
#define {{ name|upper }}_AXIS_SIZE_PRE {{ axis_size_pre }}
#endif /* {{ name|upper }}_LAYER_H */
{% filter indent(width=4, first=False) %}
{% include "./_mem_offset.jinja" %}
softmax_forward<{{ in_name[0]|upper }}_NB_CHANNELS,
{{ in_name[0]|upper }}_IN_HEIGHT,
{{ in_name[0]|upper }}_IN_WIDTH,
{{ out_name[0]|upper }}_NB_OUTPUTS,
{{ out_name[0]|upper }}_OUT_HEIGHT,
{{ out_name[0]|upper }}_OUT_WIDTH,
{{ name|upper }}_AXIS>
({{in_name[0]}}, {{out_name[0]}});
softmax_forward<{{ name|upper }}_AXIS_SIZE,
{{ name|upper }}_AXIS_SIZE_POST,
{{ name|upper }}_AXIS_SIZE_PRE>
({{in_name[0]}}, {{out_name[0]}});
{% include "./_save_outputs.jinja" %}
{% endfilter %}
......@@ -169,6 +169,30 @@ class test_operator_export(unittest.TestCase):
self.unit_test_export(model, "Softmax", [[1, 10]])
def test_export_softmax_batch(self):
print("SoftmaxBatch")
model = aidge_core.sequential([
aidge_core.Softmax(axis=1, name="sf0")
])
self.unit_test_export(model, "SoftmaxBatch", [[3, 10]])
def test_export_softmax_axis_2(self):
print("SoftmaxAxis2")
model = aidge_core.sequential([
aidge_core.Softmax(axis=2, name="sf0")
])
self.unit_test_export(model, "SoftmaxAxis2", [[1, 10, 3, 7]])
def test_export_softmax_axis_0(self):
print("SoftmaxAxis0")
model = aidge_core.sequential([
aidge_core.Softmax(axis=0, name="sf0")
])
self.unit_test_export(model, "SoftmaxAxis0", [[10]])
@unittest.skip("Currently this test is failing")
def test_export_FC_image_in(self):
"""Test exporting a FC operator with a HWC input.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment