Export Models to onnx?

[andife]

Hello, I'm new to ludwig.
Can some provide a code fragment, how to export the final model to onnx?
Thank you
Andreas

[brightsparc]

Hi @andife, yes its possible to export the predictor pytorch model trained ludwig to a static onnx graph.

1. To get started, see the following code that uses the titanic dataset

from ludwig.api import LudwigModel
from ludwig.data.preprocessing import preprocess_for_prediction
from ludwig.datasets import titanic
import torch

# Get the training and test set for titanic
training_set, test_set, _ = titanic.load(split=True)

# Load and train model config
model = LudwigModel(config="./examples/titanic/model1_config.yaml")
(
    train_stats,  # dictionary containing training statistics
    preprocessed_data,  # tuple Ludwig Dataset objects of pre-processed training data
    output_directory,  # location of training results stored on disk
) = model.train(
    training_set=training_set,
    test_set=test_set
)

# Split the preprocessed data, to extract the training set metadata.
(
     proc_training_set,
     proc_validation_set,
     proc_test_set,
     training_set_metadata
) = preprocessed_data

# Pre process the test dataset
batch_size = 1
pred_set, _ = preprocess_for_prediction(
    model.config,
    dataset=test_set.head(batch_size),
    training_set_metadata=training_set_metadata,
    include_outputs=False,
)
torch_data_to_predict = {
    name: torch.from_numpy(pred_set.dataset[feature.proc_column])
    for name, feature in model.model.input_features.items()
}

# Export the onnx model passing the model predictor, and pre-processed dataset
cols = ["Pclass", "Sex", "Age", "SibSp", "Parch", "Fare", "Embarked"]
target = "Survived"
model_path = "export.onnx"
torch.onnx.export(model.model, 
                  torch_data_to_predict, 
                  model_path, verbose=True, 
                  input_names=cols, 
                  output_names=[target])

2. Then we can onnx library to load and print the graph. First install onnx

conda install -c conda-forge onnx

Then import the library and load from the exported model path.

import onnx

onnx_model = onnx.load(model_path)
onnx.checker.check_model(onnx_model)
print(onnx.helper.printable_graph(onnx_model.graph))

3. Finally you can install the onnx runtime:

conda install --yes onnx onnxruntime

To then perform inference on the onnx graph you exported, passing pre-processed inputs in numpy format:

import onnxruntime as ort

np_data_to_predict = {
    name: pred_set.dataset[feature.proc_column]
    for name, feature in model.model.input_features.items()
}

ort_session = ort.InferenceSession(model_path)
outputs = ort_session.run(None, np_data_to_predict)

print("Survived last hidden", outputs[0])
print("Survived logits", outputs[1])

[andife]

Thank you for your fast and detailed answer.

Unfortunatly the code, does not run for me so far. The error I get is
" assert list(inputs.keys()) == self.input_features.keys() AttributeError: 'NoneType' object has no attribute 'keys'"

The complete error message:

 (ludwigai) C:\dev\testludwig>python ludwig_onnx_export.py
NumExpr defaulting to 8 threads.
C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torchaudio\backend\utils.py:62: UserWarning: No audio backend is available.
  warnings.warn("No audio backend is available.")
Torch version: 1.12.1+cpu
Ludwig version: 0.5.5
   PassengerId  Survived  Pclass                                               Name     Sex   Age  SibSp  Parch            Ticket     Fare Cabin Embarked
0            1       0.0       3                            Braund, Mr. Owen Harris    male  22.0      1      0         A/5 21171   7.2500   NaN        S
1            2       1.0       1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1      0          PC 17599  71.2833   C85        C
2            3       1.0       3                             Heikkinen, Miss. Laina  female  26.0      0      0  STON/O2. 3101282   7.9250   NaN        S
3            4       1.0       1       Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1      0            113803  53.1000  C123        S
4            5       0.0       3                           Allen, Mr. William Henry    male  35.0      0      0            373450   8.0500   NaN        S
ray.init() failed: Could not find any running Ray instance. Please specify the one to connect to by setting `--address` flag or `RAY_ADDRESS` environment variable.
---
{'Pclass': tensor([1], dtype=torch.int8), 'Sex': tensor([1], dtype=torch.int8), 'Age': tensor([34.5000]), 'SibSp': tensor([0.]), 'Parch': tensor([0.]), 'Fare': tensor([7.8292]), 'Embarked': tensor([3], dtype=torch.int8)}
----
ECD(
  (input_features): LudwigFeatureDict(
    (module_dict): ModuleDict(
      (Pclass__ludwig): CategoryInputFeature(
        (encoder_obj): CategoricalEmbedEncoder(
          (embed): Embed(
            (embeddings): Embedding(4, 4)
          )
        )
      )
      (Sex__ludwig): CategoryInputFeature(
        (encoder_obj): CategoricalEmbedEncoder(
          (embed): Embed(
            (embeddings): Embedding(3, 3)
          )
        )
      )
      (Age__ludwig): NumberInputFeature(
        (encoder_obj): PassthroughEncoder()
      )
      (SibSp__ludwig): NumberInputFeature(
        (encoder_obj): PassthroughEncoder()
      )
      (Parch__ludwig): NumberInputFeature(
        (encoder_obj): PassthroughEncoder()
      )
      (Fare__ludwig): NumberInputFeature(
        (encoder_obj): PassthroughEncoder()
      )
      (Embarked__ludwig): CategoryInputFeature(
        (encoder_obj): CategoricalEmbedEncoder(
          (embed): Embed(
            (embeddings): Embedding(5, 5)
          )
        )
      )
    )
  )
  (combiner): ConcatCombiner(
    (input_features): LudwigFeatureDict(
      (module_dict): ModuleDict(
        (Pclass__ludwig): CategoryInputFeature(
          (encoder_obj): CategoricalEmbedEncoder(
            (embed): Embed(
              (embeddings): Embedding(4, 4)
            )
          )
        )
        (Sex__ludwig): CategoryInputFeature(
          (encoder_obj): CategoricalEmbedEncoder(
            (embed): Embed(
              (embeddings): Embedding(3, 3)
            )
          )
        )
        (Age__ludwig): NumberInputFeature(
          (encoder_obj): PassthroughEncoder()
        )
        (SibSp__ludwig): NumberInputFeature(
          (encoder_obj): PassthroughEncoder()
        )
        (Parch__ludwig): NumberInputFeature(
          (encoder_obj): PassthroughEncoder()
        )
        (Fare__ludwig): NumberInputFeature(
          (encoder_obj): PassthroughEncoder()
        )
        (Embarked__ludwig): CategoryInputFeature(
          (encoder_obj): CategoricalEmbedEncoder(
            (embed): Embed(
              (embeddings): Embedding(5, 5)
            )
          )
        )
      )
    )
    (fc_stack): FCStack(
      (stack): ModuleList()
    )
  )
  (output_features): LudwigFeatureDict(
    (module_dict): ModuleDict(
      (Survived__ludwig): BinaryOutputFeature(
        (fc_stack): FCStack(
          (stack): ModuleList()
        )
        (_prediction_module): _BinaryPredict()
        (reduce_sequence_input): SequenceReducer(
          (_reduce_obj): ReduceSum()
        )
        (decoder_obj): Regressor(
          (dense): Dense(
            (dense): Linear(in_features=16, out_features=1, bias=True)
          )
        )
        (train_loss_function): BWCEWLoss(
          (loss_fn): BCEWithLogitsLoss()
        )
        (eval_loss_metric): BWCEWLMetric(
          (avg): MeanMetric()
          (loss_function): BWCEWLoss(
            (loss_fn): BCEWithLogitsLoss()
          )
        )
      )
    )
  )
  (eval_loss_metric): MeanMetric()
  (eval_additional_losses_metrics): MeanMetric()
)
-----
Traceback (most recent call last):
  File "ludwig_onnx_export.py", line 56, in <module>
    torch.onnx.export(model.model,
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\__init__.py", line 350, in export
    return utils.export(
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\utils.py", line 163, in export
    _export(
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\utils.py", line 1074, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\utils.py", line 727, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\utils.py", line 602, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\onnx\utils.py", line 517, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\jit\_trace.py", line 1175, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\jit\_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\jit\_trace.py", line 118, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\torch\nn\modules\module.py", line 1118, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "C:\Users\andife\AppData\Local\miniforge3\envs\ludwigai\lib\site-packages\ludwig\models\ecd.py", line 180, in forward
    assert list(inputs.keys()) == self.input_features.keys()
AttributeError: 'NoneType' object has no attribute 'keys'

(ludwigai) C:\dev\testludwig>

The code I used is:

from ludwig.api import LudwigModel
from ludwig.data.preprocessing import preprocess_for_prediction
from ludwig.datasets import titanic
import torch
import ludwig
import torch.onnx

print('Torch version: ' + str(torch.__version__))
print('Ludwig version: ' + str(ludwig.__version__))
# Get the training and test set for titanic
training_set, test_set, _ = titanic.load(split=True)

print(training_set.head())
# Load and train model config
model = LudwigModel(config="C:\dev\model1_config.yaml")
(
    train_stats,  # dictionary containing training statistics
    preprocessed_data,  # tuple Ludwig Dataset objects of pre-processed training data
    output_directory,  # location of training results stored on disk
) = model.train(
    training_set=training_set,
    test_set=test_set
)

# Split the preprocessed data, to extract the training set metadata.
(
     proc_training_set,
     proc_validation_set,
     proc_test_set,
     training_set_metadata
) = preprocessed_data

# Pre process the test dataset
batch_size = 1
pred_set, _ = preprocess_for_prediction(
    model.config,
    dataset=test_set.head(batch_size),
    training_set_metadata=training_set_metadata,
    include_outputs=False,
)
torch_data_to_predict = {
    name: torch.from_numpy(pred_set.dataset[feature.proc_column])
    for name, feature in model.model.input_features.items()
}

print('---')
print(torch_data_to_predict)
print('----')
print(model.model)
print('-----')

# Export the onnx model passing the model predictor, and pre-processed dataset
cols = ["Pclass", "Sex", "Age", "SibSp", "Parch", "Fare", "Embarked"]
target = "Survived"
model_path = "C:\dev\export.onnx"
torch.onnx.export(model.model, 
                  torch_data_to_predict, 
                  model_path, verbose=True, 
                  input_names=cols, 
                  output_names=[target])

[brightsparc]

Hi, it seems like there is an issue with ray backend which is the default for training.

You should be able to install this with

pip install "ludwig[distributed]"

But if you are getting errors due to windows or other reasons, you can update your config to add a backend of type local as show below, which should allow you to get up and running with that script.

input_features:
    - name: Pclass
      type: category
    - name: Sex
      type: category
    - name: Age
      type: number
      preprocessing:
          missing_value_strategy: fill_with_mean
    - name: SibSp
      type: number
    - name: Parch
      type: number
    - name: Fare
      type: number
      preprocessing:
          missing_value_strategy: fill_with_mean
    - name: Embarked
      type: category

output_features:
    - name: Survived
      type: binary
        
trainer:
    epochs: 10
        
backend:
  type: local 

[andife]

I now use directly ubuntu and the backend-type local, I currently do not really get on.
I can observe that the training generally works with the use environment, but as noted the problems with the export.

(ludwigai) andi@erde:~/Andreas/dev/ludwig/examples/titanic$ python simple_titanic.py 
NumExpr defaulting to 4 threads.
Extension horovod.torch has not been built: /home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/horovod/torch/mpi_lib/_mpi_lib.cpython-38-x86_64-linux-gnu.so not found
If this is not expected, reinstall Horovod with HOROVOD_WITH_PYTORCH=1 to debug the build error.
Warning! MPI libs are missing, but python applications are still available.
Torch version: 1.12.1+cu102
Ludwig version: 0.6.dev
   PassengerId  Survived  Pclass                                               Name     Sex   Age  SibSp  Parch            Ticket     Fare Cabin Embarked
0            1       0.0       3                            Braund, Mr. Owen Harris    male  22.0      1      0         A/5 21171   7.2500   NaN        S
1            2       1.0       1  Cumings, Mrs. John Bradley (Florence Briggs Th...  female  38.0      1      0          PC 17599  71.2833   C85        C
2            3       1.0       3                             Heikkinen, Miss. Laina  female  26.0      0      0  STON/O2. 3101282   7.9250   NaN        S
3            4       1.0       1       Futrelle, Mrs. Jacques Heath (Lily May Peel)  female  35.0      1      0            113803  53.1000  C123        S
4            5       0.0       3                           Allen, Mr. William Henry    male  35.0      0      0            373450   8.0500   NaN        S
Using training dataframe
Building dataset (it may take a while)
Building dataset: DONE
/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/sklearn/experimental/enable_hist_gradient_boosting.py:16: UserWarning: Since version 1.0, it is not needed to import enable_hist_gradient_boosting anymore. HistGradientBoostingClassifier and HistGradientBoostingRegressor are now stable and can be normally imported from sklearn.ensemble.
  warnings.warn(
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: loss. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: roc_auc. Exception: module 'horovod.torch' has no attribute 'is_initialized'
Caught exception computing metric: accuracy. Exception: module 'horovod.torch' has no attribute 'is_initialized'
---
----
-----
Traceback (most recent call last):
  File "simple_titanic.py", line 56, in <module>
    torch.onnx.export(model.model, 
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/__init__.py", line 350, in export
    return utils.export(
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/utils.py", line 163, in export
    _export(
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/utils.py", line 1074, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/utils.py", line 727, in _model_to_graph
    graph, params, torch_out, module = _create_jit_graph(model, args)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/utils.py", line 602, in _create_jit_graph
    graph, torch_out = _trace_and_get_graph_from_model(model, args)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/onnx/utils.py", line 517, in _trace_and_get_graph_from_model
    trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph(
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/jit/_trace.py", line 1175, in _get_trace_graph
    outs = ONNXTracedModule(f, strict, _force_outplace, return_inputs, _return_inputs_states)(*args, **kwargs)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/jit/_trace.py", line 127, in forward
    graph, out = torch._C._create_graph_by_tracing(
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/jit/_trace.py", line 118, in wrapper
    outs.append(self.inner(*trace_inputs))
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1118, in _slow_forward
    result = self.forward(*input, **kwargs)
  File "/home/andi/miniforge-pypy3/envs/ludwigai/lib/python3.8/site-packages/ludwig/models/ecd.py", line 145, in forward
    assert list(inputs.keys()) == self.input_features.keys()
AttributeError: 'NoneType' object has no attribute 'keys'

The model saved in for example "results/api_experiment_run_1/model$" is a LudwigModel? not converted to torchscript?
How can I load it and convert to a regular pytorch-model? (Maybe than I could get help of the pytorch community?)

Furthermore, I am even more interested in the export of an image classification problem. (so mnist might an even better example for me) But there I currently get the same error. Has anyone else observed this with den keys?

[brightsparc]

Hi @andife please make sure your config is well formed yaml. It seems like the training process is not able to start due to the inputs not being recognised in your model.