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.