Skip to content

Commit 8c70508

Browse files
gluonparticleguitargeek
gluonparticle
authored and
guitargeek
committed
[tmva][sofie] Add InstanceNormalization operator for ONNX inference
Add ROperator_InstanceNormalization, implementing the ONNX InstanceNormalization operator: for every sample and channel of an input of shape (N, C, D1, ..., Dn), the elements over the spatial dimensions are normalized to zero mean and unit variance, then scaled and shifted with the per-channel scale and B tensors. The operator is registered in RModelParser_ONNX through a new ParseInstanceNormalization.cxx, following the one-file-per-operator convention of the other operators. Only static shapes and float inputs are supported; other input types are rejected with an error, as in ParseLayerNormalization. Three models are added to generate_input_models.py, with their expected outputs computed by onnx's ReferenceEvaluator: * InstanceNormalization, of shape (2, 3, 4, 5) with asymmetric scale and bias, so that a wrong per-instance offset or a scale indexed on the wrong axis does not cancel out, * InstanceNormalization3d, of rank 3, where the spatial size is not a product of several dimensions, * InstanceNormalizationEpsilon, with a non-default epsilon of 0.01 and a small-variance input, so that the attribute dominates the result. Two mutations of the operator were checked to be caught by these tests: ignoring the epsilon attribute fails InstanceNormalizationEpsilon, and indexing the scale by the wrong axis fails the other two. 🤖 Done with the help of AI.
1 parent 86cda43 commit 8c70508

6 files changed

Lines changed: 292 additions & 0 deletions

File tree

‎tmva/sofie/inc/TMVA/ROperator_InstanceNormalization.hxx‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
#ifndef TMVA_SOFIE_ROPERATOR_INSTANCENORMALIZATION
2+
#define TMVA_SOFIE_ROPERATOR_INSTANCENORMALIZATION
3+
4+
#include "TMVA/ROperator.hxx"
5+
#include "TMVA/RModel.hxx"
6+
#include "TMVA/SOFIE_common.hxx"
7+
8+
#include <sstream>
9+
#include <stdexcept>
10+
11+
namespace TMVA {
12+
namespace Experimental {
13+
namespace SOFIE {
14+
15+
/*! \class ROperator_InstanceNormalization
16+
\brief Implementation of the ONNX InstanceNormalization operator.
17+
18+
The input X has shape (N, C, D1, ..., Dn). For every sample n and channel c,
19+
the elements over the spatial dimensions D1, ..., Dn are normalized to zero
20+
mean and unit variance, and then scaled and shifted with the per-channel
21+
scale and B tensors, which both have shape (C):
22+
23+
Y[n, c, ...] = scale[c] * (X[n, c, ...] - mean[n, c]) / sqrt(var[n, c] + epsilon) + B[c]
24+
*/
25+
template <typename T>
26+
class ROperator_InstanceNormalization final : public ROperator {
27+
private:
28+
float fEpsilon;
29+
std::string fNInput;
30+
std::string fNScale;
31+
std::string fNBias;
32+
std::string fNOutput;
33+
std::vector<size_t> fShape;
34+
std::string fType;
35+
36+
public:
37+
ROperator_InstanceNormalization() : fEpsilon(1e-5) {}
38+
39+
ROperator_InstanceNormalization(float epsilon, std::string nameInput, std::string nameScale, std::string nameBias,
40+
std::string nameOutput)
41+
: fEpsilon(epsilon),
42+
fNInput(UTILITY::Clean_name(nameInput)),
43+
fNScale(UTILITY::Clean_name(nameScale)),
44+
fNBias(UTILITY::Clean_name(nameBias)),
45+
fNOutput(UTILITY::Clean_name(nameOutput))
46+
{
47+
fInputTensorNames = {fNInput, fNScale, fNBias};
48+
fOutputTensorNames = {fNOutput};
49+
}
50+
51+
std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> inputShapes) override
52+
{
53+
return {inputShapes[0]};
54+
}
55+
56+
void Initialize(RModel &model) override
57+
{
58+
for (const std::string &name : {fNInput, fNScale, fNBias}) {
59+
if (!model.CheckIfTensorAlreadyExist(name)) {
60+
throw std::runtime_error("TMVA::SOFIE - InstanceNormalization - Tensor " + name + " not found.");
61+
}
62+
}
63+
64+
fShape = model.GetTensorShape(fNInput);
65+
if (fShape.size() < 3) {
66+
throw std::runtime_error("TMVA::SOFIE - InstanceNormalization - Input tensor " + fNInput + " has rank " +
67+
std::to_string(fShape.size()) + " but at least rank 3 (N, C, D1, ...) is required.");
68+
}
69+
70+
// scale and B are 1D tensors of length C
71+
const size_t channels = fShape[1];
72+
for (const std::string &name : {fNScale, fNBias}) {
73+
auto shape = model.GetTensorShape(name);
74+
if (shape.size() != 1 || shape[0] != channels) {
75+
throw std::runtime_error("TMVA::SOFIE - InstanceNormalization - Tensor " + name + " has invalid shape " +
76+
ConvertShapeToString(shape) + ", expected " +
77+
ConvertShapeToString(std::vector<size_t>{channels}) + ".");
78+
}
79+
}
80+
81+
fType = ConvertTypeToString(model.GetTensorType(fNInput));
82+
model.AddIntermediateTensor(fNOutput, model.GetTensorType(fNInput), fShape);
83+
}
84+
85+
std::string Generate(std::string) override
86+
{
87+
if (fShape.empty()) {
88+
throw std::runtime_error("TMVA::SOFIE - InstanceNormalization called to generate without being initialized.");
89+
}
90+
91+
const size_t batchSize = fShape[0];
92+
const size_t channels = fShape[1];
93+
size_t spatialSize = 1;
94+
for (size_t i = 2; i < fShape.size(); ++i)
95+
spatialSize *= fShape[i];
96+
97+
std::stringstream out;
98+
out << "\n" << SP << "//---- InstanceNormalization " << fNOutput << "\n";
99+
out << SP << "for (size_t n = 0; n < " << batchSize << "; n++) {\n";
100+
out << SP << SP << "for (size_t c = 0; c < " << channels << "; c++) {\n";
101+
out << SP << SP << SP << "const size_t offset = n * " << channels * spatialSize << " + c * " << spatialSize
102+
<< ";\n";
103+
104+
// Compute the mean over the spatial dimensions
105+
out << SP << SP << SP << fType << " mean = 0.;\n";
106+
out << SP << SP << SP << "for (size_t i = 0; i < " << spatialSize << "; i++) {\n";
107+
out << SP << SP << SP << SP << "mean += tensor_" << fNInput << "[offset + i];\n";
108+
out << SP << SP << SP << "}\n";
109+
out << SP << SP << SP << "mean /= " << fType << "(" << spatialSize << ");\n";
110+
111+
// Compute the inverse standard deviation from the deviations around the mean
112+
out << SP << SP << SP << fType << " sum = 0.;\n";
113+
out << SP << SP << SP << "for (size_t i = 0; i < " << spatialSize << "; i++) {\n";
114+
out << SP << SP << SP << SP << fType << " tmp = tensor_" << fNInput << "[offset + i] - mean;\n";
115+
out << SP << SP << SP << SP << "sum += tmp * tmp;\n";
116+
out << SP << SP << SP << "}\n";
117+
out << SP << SP << SP << fType << " invStdDev = 1 / std::sqrt(sum / " << fType << "(" << spatialSize << ") + "
118+
<< std::to_string(fEpsilon) << ");\n";
119+
120+
// Y = scale o invStdDev (X - mean) + B
121+
out << SP << SP << SP << fType << " scale = tensor_" << fNScale << "[c];\n";
122+
out << SP << SP << SP << fType << " bias = tensor_" << fNBias << "[c];\n";
123+
out << SP << SP << SP << "for (size_t i = 0; i < " << spatialSize << "; i++) {\n";
124+
out << SP << SP << SP << SP << "tensor_" << fNOutput << "[offset + i] = scale * (tensor_" << fNInput
125+
<< "[offset + i] - mean) * invStdDev + bias;\n";
126+
out << SP << SP << SP << "}\n";
127+
128+
out << SP << SP << "}\n";
129+
out << SP << "}\n";
130+
return out.str();
131+
}
132+
};
133+
134+
} // namespace SOFIE
135+
} // namespace Experimental
136+
} // namespace TMVA
137+
#endif

‎tmva/sofie/test/TestCustomModelsFromONNX.cxx‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -948,6 +948,33 @@ TEST(ONNX, LayerNormalization4d)
948948
expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
949949
}
950950

951+
TEST(ONNX, InstanceNormalization)
952+
{
953+
SofieReference ref = readReference("InstanceNormalization");
954+
955+
ASSERT_INCLUDE_AND_RUN(std::vector<float>, "InstanceNormalization", ref.f32("input0"));
956+
957+
expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
958+
}
959+
960+
TEST(ONNX, InstanceNormalization3d)
961+
{
962+
SofieReference ref = readReference("InstanceNormalization3d");
963+
964+
ASSERT_INCLUDE_AND_RUN(std::vector<float>, "InstanceNormalization3d", ref.f32("input0"));
965+
966+
expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
967+
}
968+
969+
TEST(ONNX, InstanceNormalizationEpsilon)
970+
{
971+
SofieReference ref = readReference("InstanceNormalizationEpsilon");
972+
973+
ASSERT_INCLUDE_AND_RUN(std::vector<float>, "InstanceNormalizationEpsilon", ref.f32("input0"));
974+
975+
expectNear(output, ref.f32("output0"), DEFAULT_TOLERANCE);
976+
}
977+
951978
TEST(ONNX, Equal)
952979
{
953980
SofieReference ref = readReference("Equal");

‎tmva/sofie/test/generate_input_models.py‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1642,6 +1642,61 @@ def make_HardSwish():
16421642
return _model(graph, opset=14, ir_version=13)
16431643

16441644

1645+
def _instance_normalization(shape, scale, bias, **kwargs):
1646+
"""InstanceNormalization graph over an input of the given shape. scale and
1647+
B are per-channel initializers of length shape[1]."""
1648+
channels = shape[1]
1649+
nodes = [
1650+
helper.make_node('InstanceNormalization', ['X', 'scale', 'B'], ['Y'], **kwargs),
1651+
]
1652+
graph = helper.make_graph(
1653+
nodes,
1654+
'InstanceNormalization',
1655+
inputs=[
1656+
_vi('X', FLOAT, shape),
1657+
_vi('scale', FLOAT, [channels]),
1658+
_vi('B', FLOAT, [channels]),
1659+
],
1660+
outputs=[
1661+
_vi('Y', FLOAT, shape),
1662+
],
1663+
initializer=[
1664+
_tensor('scale', FLOAT, [channels], scale),
1665+
_tensor('B', FLOAT, [channels], bias),
1666+
],
1667+
)
1668+
return _model(graph, opset=17, ir_version=8)
1669+
1670+
1671+
def make_InstanceNormalization():
1672+
"""Ops: InstanceNormalization"""
1673+
# Batch and channel size both > 1, so that a wrong per-instance offset or a
1674+
# scale/bias indexed by the wrong axis does not go unnoticed.
1675+
return _instance_normalization([2, 3, 4, 5],
1676+
scale=[0.5, 1.0, -2.0],
1677+
bias=[0.0, 0.25, -1.5])
1678+
1679+
1680+
def make_InstanceNormalization3d():
1681+
"""Ops: InstanceNormalization"""
1682+
# Rank 3 (N, C, D): exercises a spatial size that is not a product of
1683+
# several dimensions.
1684+
return _instance_normalization([2, 2, 6],
1685+
scale=[1.5, -0.5],
1686+
bias=[1.0, 2.0])
1687+
1688+
1689+
def make_InstanceNormalizationEpsilon():
1690+
"""Ops: InstanceNormalization"""
1691+
# Non-default epsilon, combined with the small-variance input in
1692+
# TEST_INPUTS, so that the attribute dominates the result and a parser that
1693+
# ignored it would be caught.
1694+
return _instance_normalization([1, 2, 3, 3],
1695+
scale=[1.0, 1.0],
1696+
bias=[0.0, 0.0],
1697+
epsilon=0.01)
1698+
1699+
16451700
def make_IsInf():
16461701
"""Ops: IsInf"""
16471702
nodes = [
@@ -4741,6 +4796,9 @@ def make_Where():
47414796
'GreaterOrEqual': make_GreaterOrEqual,
47424797
'HardSigmoid': make_HardSigmoid,
47434798
'HardSwish': make_HardSwish,
4799+
'InstanceNormalization': make_InstanceNormalization,
4800+
'InstanceNormalization3d': make_InstanceNormalization3d,
4801+
'InstanceNormalizationEpsilon': make_InstanceNormalizationEpsilon,
47444802
'IsInf': make_IsInf,
47454803
'LSTMBatchwise': make_LSTMBatchwise,
47464804
'LSTMBidirectional': make_LSTMBidirectional,
@@ -4940,6 +4998,21 @@ def rand_f32(seed, shape):
49404998
],
49414999
'HardSigmoid': [f32([1.0, -2.0, 3.0, 0.5, -1.0, 2.0], (6,))],
49425000
'HardSwish': [f32([1.0, -2.0, 3.0, 0.5, -1.0, 2.0], (6,))],
5001+
# Per (n, c) slice a different mean and variance, so that a normalization
5002+
# that mixed up instances or channels would not cancel out.
5003+
'InstanceNormalization': [
5004+
f32(np.arange(120.0) * 0.1 - 6.0 + (np.arange(120.0) % 7), (2, 3, 4, 5)),
5005+
],
5006+
'InstanceNormalization3d': [
5007+
f32([1.0, 2.0, 3.0, 4.0, 5.0, 6.0,
5008+
-1.0, -3.0, 0.0, 2.0, -2.0, 4.0,
5009+
10.0, 10.0, 10.0, 10.0, 10.0, 11.0,
5010+
0.5, -0.5, 1.5, -1.5, 2.5, -2.5], (2, 2, 6)),
5011+
],
5012+
# Small variance (~2e-3), comparable to the model's epsilon of 0.01.
5013+
'InstanceNormalizationEpsilon': [
5014+
f32((np.arange(18.0) % 5 - 2.0) * 0.025, (1, 2, 3, 3)),
5015+
],
49435016
'LSTMBatchwise': [f32(np.arange(1.0, 7.0), (3, 1, 2))],
49445017
'LSTMBidirectional': [f32(np.arange(1.0, 7.0), (3, 1, 2))],
49455018
'LSTMDefaults': [f32(np.arange(1.0, 7.0), (3, 1, 2))],

‎tmva/sofie_parsers/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofieParser
5353
src/ParseErf.cxx
5454
src/ParseRange.cxx
5555
src/ParseLayerNormalization.cxx
56+
src/ParseInstanceNormalization.cxx
5657
src/ParseExpand.cxx
5758
src/ParseGather.cxx
5859
src/ParseGatherND.cxx

‎tmva/sofie_parsers/src/ParseInstanceNormalization.cxx‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
#include "TMVA/RModelParser_ONNX.hxx"
2+
#include "TMVA/ROperator_InstanceNormalization.hxx"
3+
#include "onnx_proto3.pb.h"
4+
5+
namespace TMVA {
6+
namespace Experimental {
7+
namespace SOFIE {
8+
9+
ParserFuncSignature ParseInstanceNormalization = [](RModelParser_ONNX &parser,
10+
const onnx::NodeProto &nodeproto) -> std::unique_ptr<ROperator> {
11+
ETensorType input_type = ETensorType::UNDEFINED;
12+
const std::string input_name = nodeproto.input(0);
13+
if (parser.IsRegisteredTensorType(input_name)) {
14+
input_type = parser.GetTensorType(input_name);
15+
} else {
16+
throw std::runtime_error("TMVA::SOFIE ONNX Parser InstanceNormalization op has input tensor " + input_name +
17+
" but its type is not yet registered");
18+
}
19+
20+
float epsilon = 1e-5;
21+
for (int64_t i = 0; i < nodeproto.attribute_size(); i++) {
22+
if (nodeproto.attribute(i).name() == "epsilon") {
23+
epsilon = nodeproto.attribute(i).f();
24+
}
25+
}
26+
27+
// Inputs: X (0), scale (1), B (2)
28+
const std::string name_scale = nodeproto.input(1);
29+
const std::string name_bias = nodeproto.input(2);
30+
const std::string output_name = nodeproto.output(0);
31+
32+
std::unique_ptr<ROperator> op;
33+
switch (input_type) {
34+
case ETensorType::FLOAT:
35+
op.reset(new ROperator_InstanceNormalization<float>(epsilon, input_name, name_scale, name_bias, output_name));
36+
break;
37+
default:
38+
throw std::runtime_error("TMVA::SOFIE ONNX parser Operator with input type " + ConvertTypeToString(input_type) +
39+
" not supported.");
40+
break;
41+
}
42+
43+
if (!parser.IsRegisteredTensorType(output_name)) {
44+
parser.RegisterTensorType(output_name, input_type);
45+
}
46+
47+
return op;
48+
};
49+
50+
} // namespace SOFIE
51+
} // namespace Experimental
52+
} // namespace TMVA

‎tmva/sofie_parsers/src/RModelParser_ONNX.cxx‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,7 @@ extern ParserFuncSignature ParseExpand;
8484
extern ParserFuncSignature ParseShape;
8585
extern ParserFuncSignature ParseMatMul;
8686
extern ParserFuncSignature ParseLayerNormalization;
87+
extern ParserFuncSignature ParseInstanceNormalization;
8788
extern ParserFuncSignature ParseGather;
8889
extern ParserFuncSignature ParseGatherND;
8990
extern ParserFuncSignature ParseErf;
@@ -377,6 +378,7 @@ RModelParser_ONNX::RModelParser_ONNX() noexcept : fOperatorsMapImpl(std::make_un
377378
RegisterOperator("Tile", ParseTile);
378379
RegisterOperator("Split", ParseSplit);
379380
RegisterOperator("If", ParseIf);
381+
RegisterOperator("InstanceNormalization", ParseInstanceNormalization);
380382
RegisterOperator("Pad", ParsePad);
381383
RegisterOperator("Where", ParseWhere);
382384
RegisterOperator("Einsum", ParseEinsum);

0 commit comments

Comments
 (0)