How to Validate Your ONNX Models in Python: Briefly
Before exporting, you need to have the model and input data ready. It's usually a simple matter of instantiating the model and preparing the data. For example, if you are working with a simple neural network, it might look like this:
import torch
import torchvision.models as models
# Загружаем предобученную модель ResNet
model = models.resnet18(pretrained=True)
model.eval() # Переключаем модель в режим оценки
# Создаем случайный тензор для входных данных
dummy_input = torch.randn(1, 3, 224, 224) # Формат: [batch_size, channels, height, width]Now that the model is ready, it's time to export it. For this we use the method torch.onnx.export():
torch.onnx.export(model,
dummy_input,
"model.onnx",
export_params=True,
opset_version=11, # Версия ONNX
do_constant_folding=True, # Оптимизация
input_names=['input'], # Названия входных слоев
output_names=['output']) # Названия выходных слоев
print("Модель успешно экспортирована в формат ONNX!")A couple of tips
The model must be in evaluation mode model.eval()to turn off layers like Dropout And BatchNormwhich may lead to instability during testing.
If the model uses custom operations, they must be supported in ONNX, otherwise an error may occur during export.
Now let's look at the export process from TensorFlow. Let's start with a simple example using Keras.
Let's create a simple model using Keras:
import tensorflow as tf
# Создаем простую модель
model = tf.keras.Sequential([
tf.keras.layers.Dense(128, activation='relu', input_shape=(784,)),
tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy", metrics=['accuracy'])
# Создаем случайный тензор для входных данных
dummy_input = tf.random.normal([1, 784]) # Формат: [batch_size, features]To export a model from TensorFlow to ONNX, we use the library tf2onnxwhich allows you to convert TensorFlow models to ONNX format:
pip install tf2onnxNow we can export the model:
import tf2onnx
# Экспортируем модель
onnx_model = tf2onnx.convert.from_keras(model, output_path="model.onnx")
print("Модель успешно экспортирована в формат ONNX!")Testing ONNX models
Before running tests, you need to decide on the test data. It is important to consider several factors here:
Data type: If you are working with images, use the datasets that were used to train the model. For text data, variety is needed so that the model can handle different cases.
Test set size: The test set must be large to cover all possible scenarios. The more data, the higher the likelihood of identifying problems.
Evaluation criteria should include performance metrics such as:
Accuracy: Compare your model's predictions with actual values. Use metrics: precision, recall and F1-measure.
Lead time: Measure the time required to make predictions on test data. You can use Python's built-in time functions, like this
time.time().Memory usage: Check how well the model uses resources. You can use the library
memory-profilerto track your model's memory usage at runtime.
There are several tools on the market for testing ONNX models, and we will focus on the two main ones:
ONNX Runtime: This is a high-performance version of ONNX models that provides fast and efficient operation.
NumPy: For analyzing and processing data, as well as for comparing predictions with actual values.
Now let's move on to practice.
pip install onnxruntime numpyLet's start by loading our model and preparing test data. Let's say there is a model that takes 224×224 pixel images as input:
import onnx
import onnxruntime as ort
import numpy as np
from PIL import Image
# Загрузка модели
onnx_model = onnx.load("model.onnx")
ort_session = ort.InferenceSession("model.onnx")
# Функция для подготовки изображения
def preprocess_image(image_path):
image = Image.open(image_path)
image = image.resize((224, 224))
image = np.array(image).astype(np.float32) # Преобразуем в float32
image = image.transpose(2, 0, 1) # Изменяем порядок осей на [C, H, W]
image = np.expand_dims(image, axis=0) # Добавляем размерность [1, C, H, W]
return image
# Пример загрузки изображения
test_image = preprocess_image("test_image.jpg")Now you can make a prediction based on the prepared image:
# Выполняем предсказание
outputs = ort_session.run(None, {onnx_model.graph.input[0].name: test_image})
# Вывод предсказания
predicted_class = np.argmax(outputs[0]) # Получаем класс с максимальной вероятностью
print(f"Предсказанный класс: {predicted_class}")Now let's test how quickly our model works on several images. Let's create a loop for testing:
import time
test_images = ["image1.jpg", "image2.jpg", "image3.jpg"] # Список изображений для тестирования
total_time = 0
for image_path in test_images:
image = preprocess_image(image_path)
start_time = time.time()
outputs = ort_session.run(None, {onnx_model.graph.input[0].name: image})
end_time = time.time()
predicted_class = np.argmax(outputs[0])
total_time += end_time - start_time
print(f"Изображение: {image_path}, Предсказанный класс: {predicted_class}, Время: {end_time - start_time:.4f} секунд")
print(f"Среднее время предсказания: {total_time / len(test_images):.4f} секунд")Conclusion
Testing models in the ONNX format allows you to confidently launch your solutions into production. Don't forget that quality testing is the key to the success of your project.
To gain up-to-date knowledge on ML, come to the next open lessons:
