onnx2fmu 0.2.5
A tool for wrapping ONNX netwroks in FMUs.
Requires Python: >=3.10

# ONNX2FMU: Encapsulate ONNX models in Functional Mock-up Units (FMUs)

**What do ONNX2FMU do?** It wraps [ONNX](https://onnx.ai/) models into
co-simulation FMUs.

## 🚀 Get started

- Python 3.10+
- CMake 3.22+
- A code compiler for the host platform, which could be one of Linux, Windows
or MacOS.

The default Windows CMake Generator is Visual Studio 2022.

To install ONNX2FMU use
```bash
pip install onnx2fmu
```
in your shell.
You don't need to install CMake because it will be installed with the Python
package, but you need to install a C compiler (e.g., Visual Studio in Windows,
gcc in Linux, etc.)

## 📝 ONNX model declaration

ONNX2FMU can handle models with multiple inputs, outputs, and local variables.
These entries must be listed in the model description JSON file, and their
names must match the name of a node in the ONNX model graph.

### Model description file

A model description is declared in a JSON file and its schema includes the
following global items:

- `"name"` is the model name, which will also be the FMU archive name;
- `"description"` provides a generic description of the model;
- `"FMIVersion"` is the FMI standard version for generatign the FMU code and
the FMU binaries, which can be either `2.0` or `3.0`;
- `"inputs"` and `"outputs"` are the lists of inputs and output nodes in the
ONNX model;
- `"locals"` are mapping between an input and an output node. Their behavior is
explained in [A model with local variables](#a-model-with-local-variables).

Each entry of the the inputs and output lists is characterized by the following
schema:

- `"name"` must match the name of one of the model nodes, whereas
- `"labels"` is the list of user-provided names for each of the node elements.
The number of names in the `"labels"` list must match the number of elements of
a given entry.
- `"description"` allows the user to attach a description to each of the
arrays.

The following is an example of a model description for a model with three
input nodes and one output node.

```json
{
    "name": "example1",
    "description": "The model defines a simple example model with a scalar input and two vector inputs, one with 'local' variability and one with 'continuous' variability.",
    "FMIVersion": "2.0",
    "inputs": [
        {
            "name": "scalar_input",
            "description": "A scalar input to the model."
        },
        {
            "name": "vector_input",
            "description": "A vector of input variables with variability discrete."
        },
        {
            "name": "vector_input_discrete",
            "description": "Inputs have variability discrete by default."
        }
    ],
    "outputs": [
        {
            "name": "output",
            "description": "The output array.",
            "labels": [
                "Class1",
                "Class2",
                ...
            ]
        }
    ]
}
```

### Variability of model variables

Allowed variables types are `input`, `output`, and `local`.
Admissible variabilities are `continuous` and `discrete`;
the default choice is `continuous` if nothing is specified in the model
description.

### Model declaration: A PyTorch example

ONNX2FMU works with any ONNX model, which can be generated by all the major
ML/DL frameworks, e.g., PyTorch, TensorFlow, Scikit-Learn, etc.
However, we chose PyTorch to show how ONNX2FMU works.

The following model is used in `tests/example1` to perform some basic vector
operations.
```python
class ExampleModel(nn.Module):

    def __init__(self):
        super(ExampleModel, self).__init__()

    def forward(self, x1, x2, x3):
        # Input x1 is a scalar
        # Input x2 is a vector with causality 'local' and 4 elements
        # Input x3 is a vector with causality 'countinuous' and 5 elements
        x4 = x2 + x3[:4]
        x5 = x2 - x3[:4]
        x6 = x1 * x3[-1]
        x7 = x1 / x3[-1]
        x = torch.cat([x4, x5, x6, x7])
        return x
```
All the basic array operations are allowed, which is useful to define not only
deep learning models, but also generic, graph-based models.

A more complex examples is explained in `tests/example3`, where an RNN model
is used to predict the temperature of a point on a metallic plate.
The model is declared as follows
```python
class HeatRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers, norm_params):
        super(HeatRNN, self).__init__()
        x_min, x_max, y_min, y_max = norm_params
        self.register_buffer("x_min", torch.tensor(x_min))
        self.register_buffer("x_max", torch.tensor(x_max))
        self.register_buffer("y_min", torch.tensor(y_min))
        self.register_buffer("y_max", torch.tensor(y_max))
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, output_size)

    def forward(self, x, h=None):
        x = (x - self.x_min) / (self.x_max - self.x_min)
        out, _ = self.rnn(x, h)
        out = self.fc(out)  # Predict next time step
        out = out * (self.y_max - self.y_min) + self.y_min
        return out
```
In this example, the normalization parameters, which do not have to be
optimized, are stored in the model using the [`register_buffer`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer) method, which
detaches them from the computational graph.

### A model with local variables

Recurrent neural network architectures might require to feed the model with the
output of the model itself in a feedback loop fashion.
In ONNX2FMU, we realize this functionality through FMI `local` variables.
Mapping model's input to output is necessary though.
In [example4](tests/example4/), we show how FMUs with local variables are
declared.
The model description file must present a section named `locals` like the
following:

```json
{
    ...
    "locals": [
        {
            "nameIn": "X",
            "nameOut": "X1",
            "description": "The history of states from t-N to t."
        },
        {
            "nameIn": "U",
            "nameOut": "U1",
            "description": "The history of control variables frmo t-N to t-1."
        }
    ]
}
```

Local variables requires two names, i.e., `nameIn` and `nameOut`, which define
an input-output relationship that feeds the `nameIn` input node with the
output of the `nameOut` output node.
The user must take care to provide the right output in the `forward` method.
In `example4`, the relationship between input `U` and output `U1` is defined
as follows

```python
class ExampleModel(nn.Module):

    def __init__(self):
        super(ExampleModel, self).__init__()

    def forward(self, u, U, X):
        U1 = torch.concat((U[1:, :], u))
        x = torch.stack([U1[-3, 0], U1[-2, 1], U1[-1, 2]]).unsqueeze(0)
        X1 = torch.concat((X[1:, :], x))
        return x, X1, U1
```

In the example above, the input `U` is updated with the content of `u` and
returned by the function with the name `U1` (remember that node names cannot be
repeated in ONNX). During the next time step, the FMU will take care to pass
the output to the model as a new, updated input.

## 🔨 ONNX model generation

ONNX2FMU provides two ways to build an FMU from an ONNX model.

### CLI

ONNX2FMU is designed as command line application first. The `build` command
requires the ONNX model path and the model description path

```bash
onnx2fmu build <model.onnx> <modelDescription.json> [OPTIONS]
```

### Built-in functions

FMUs of ONNX models can be build from a Python script by calling the `build`
function from `app.py`.
This is particularly useful when training a model and generating the FMU in
the same file.

### Generation and compilation

ONNX2FMU provides the possibility to separate generation of the FMU code
and its compilation.
This can be achived using the commands `generate` and `compile`, respectively.
To see the documentation of the commands, one can use
```bash
onnx2fmu [generate|compile] --help
```
Separating FMU code generation and compilation allows a user to customize the
the FMU source code.

## Acknowledgements

The code in this repository is inspired to the [Reference FMU](https://github.com/modelica/Reference-FMUs/)
project.