Compiling Chronos-Bolt Models into ONNX

[Yashil]

Hello Chronos Community,

I'm currently working with the Chronos-Bolt models and aiming to export them to the ONNX format for deployment. However, I'm encountering issues during the export process related to unsupported operators and type errors.

Has anyone successfully exported a Chronos-Bolt model to ONNX? If so, could you share your experience or provide guidance on how to achieve this?

Any help or pointers would be greatly appreciated!

Thank you, Yashil.

[canerturkmen]

Hi @Yashil thanks for reaching out. I have a PoC here which I put together a while back. Would be great to hear if it works for you.

[Bigb20]

Hello, do you have any inference code for ONNX?

[lostella]

@Bigb20 it's linked in the answer: https://github.com/canerturkmen/chronos-forecasting/tree/onnx/scripts/onnx

[Bigb20]

What I mean is to load the ONNX model for inference of the Chronos-Bolt-Mini model and achieve the same inference effect as in the example. I want to do this without using PyTorch and also specify the prediction days and input parameters. Thank you for your reply!

import pandas as pd
import torch
from chronos import BaseChronosPipeline

pipeline = BaseChronosPipeline.from_pretrained(
"amazon/chronos-t5-small",
device_map="cuda",
torch_dtype=torch.bfloat16,
)

df = pd.read_csv(
"https://raw.githubusercontent.com/AileenNielsen/TimeSeriesAnalysisWithPython/master/data/AirPassengers.csv"
)

quantiles, mean = pipeline.predict_quantiles(
context=torch.tensor(df["#Passengers"]),
prediction_length=12,
quantile_levels=[0.1, 0.5, 0.9],
)
from chronos import ChronosPipeline, ChronosBoltPipeline

print(ChronosPipeline.predict.doc) # for Chronos models
print(ChronosBoltPipeline.predict.doc) # for Chronos-Bolt models
import matplotlib.pyplot as plt # requires: pip install matplotlib

forecast_index = range(len(df), len(df) + 12)
low, median, high = quantiles[0, :, 0], quantiles[0, :, 1], quantiles[0, :, 2]

plt.figure(figsize=(8, 4))
plt.plot(df["#Passengers"], color="royalblue", label="historical data")
plt.plot(forecast_index, median, color="tomato", label="median forecast")
plt.fill_between(forecast_index, low, high, color="tomato", alpha=0.3, label="80% prediction interval")
plt.legend()
plt.grid()
plt.show()

[lostella]

@Bigb20 the final part in the linked script only depends on onnx and numpy (no torch), you should be able to adjust it to mimic the README example

[Bigb20]

@Bigb20 the final part in the linked script only depends on onnx and numpy (no torch), you should be able to adjust it to mimic the README example

Yes, but I'm not very good at this kind of reasoning.

[quachn]

Hi @canerturkmen,
When I executed your script, I got
torch.onnx.errors.UnsupportedOperatorError: Exporting the operator 'aten::nanmean' to ONNX opset version 17 is not supported.

Thank you,
Nai Minh

[canerturkmen]

Hi @quachn , did you also check out the branch ? Note that it's the whole branch that's edited, not just the script.

[quachn]

You're right. It works fine now. Thank you @canerturkmen.

[shgoyal33]

Hi @canerturkmen I tried out the branch and it works for me for chronos bolt models. I see that the fix in your branch replaces the nan operator with a different operator which works in onnx. Now I also want to export the following checkpoint: amazon/chronos-t5-small. When I try to export this, I get the error:

\chronos-env\Lib\site-packages\transformers\generation\utils.py", line 969, in _get_logits_processor MinLengthLogitsProcessor( File "chronos-env\Lib\site-packages\transformers\generation\logits_process.py", line 134, in __init__ raise ValueError(f"min_lengthhas to be a non-negative integer, but is {min_length}") ValueError:min_length has to be a non-negative integer, but is 65

Do you know where exactly the problem lies for this and what is the fix here?

I'm using the following code for this:

import numpy as np
import onnxruntime
import torch
from chronos import ChronosBoltPipeline
from chronos import ChronosPipeline,BaseChronosPipeline

pipeline = BaseChronosPipeline.from_pretrained(
        "amazon/chronos-t5-small"
        )

batch_size, sequence_length = 32, 512

torch.onnx.export(
    pipeline.model,
    (torch.rand(batch_size, sequence_length).to(torch.int32),),
    f="amazon-chronos-t5-small.onnx")

[nai-kon]

Hi @canerturkmen. Thank you for the great PoC. It worked fine.
However the predicted length is fixed to 64 and unable to predict long sequence more than 64.
It seems that 64 is the default prediction_length of the model.
https://huggingface.co/amazon/chronos-bolt-small/blob/main/config.json#L10

Is it possible to change the prediction length?