Экспорт Julia-нейросетей в формат ONNX
Экспорт нейросети в формат ONNX
Создаём нейросеть и сохраняем ее в формате ONNX для дальнейшего использования.
Описание задачи
Чтобы развернуть готовую нейросеть в целевом окружении или даже встроить в код для микроконтроллера, распространённым решением является перевод нейросети в стандартный, независимый от фреймворков формат ONNX (Open Neural Network eXchange).
Самой распространённой библиотека для создания и обучения нейросетей в Julia называется Flux. Она позволяет создавать полносвязанные, сверточные, рекуррентные нейросети и графы других топологий, писать собственные модели нейронов и слоёв, а обучение и прямой расчет производятся в скомпилированной, бинарной форме, по скорости не уступающей языку Си.
Функции экспорта Flux-нейросетей в ONNX встречаются в нескольких библиотеках, но на данный момент наиболее успешной является инструмент экспорта, входящий в библиотеку NaiveNAXflux. Эта библиотека содержит инструменты для взаимосвязи библиотеки Flux с библиотекой для поиска оптимальных топологий нейросетей (NAS - Neural Architecture Search). NaiveNAS берет на себя работу с весами и изменение топологии нейросети, а сопряженные с ним пакеты, вроде NaiveGAflux, управляют генетической оптимизацией.
Мы также установим библиотеку ONNXRunTime, которая позволит нам запускать и тестировать сохранённую нейросеть.
Сохранение нейросети в ONNX
Установим нужные библиотеки:
Pkg.add(["Flux", "ONNXNaiveNASflux", "ONNXRunTime", "ONNXLowLevel"])
После их подключения мы сможем экспортировать в ONNX простую полносвязанную сеть:
using Flux
using ONNXNaiveNASflux
# Определим модель
model = Chain(
Dense(784 => 32, relu),
Dense(32 => 10)
)
# Экспортируем в ONNX при помощи save()
ONNXNaiveNASflux.save("fc_model.onnx", model, (784, 1));
Последним аргументом этой команды мы указываем размерность входных данных в формате количество признаков, размер батча (features, batch_size).
Количество признаков в описании объекта можно получить из свойств нейросети Flux, но размер батча во Flux не зафиксирован.
first_layer = model[1]
input_size = size(first_layer.weight, 2)
Выполнение ранее сохраненной нейросети
Запустим сохранённую нейросеть (веса в ней инициализированы случайными числами, поэтому на выходе мы получим шум):
import ONNXRunTime as ORT
fc_model = ORT.load_inference("fc_model.onnx")
fc_model
Теперь используем названия весов чтобы задать входы и получить выходы из сохраненной ONNX нейросети:
result = []
for x in eachcol(rand(Float32, 784, 20))
input_data = Dict("data_0" => reshape(x, 1, :))
outputs = fc_model(input_data)
append!(result, [outputs["dense_1"]])
end
gr()
plot(vcat(result...), title="Прогноз нейросети для случайных объектов")
Экспорт других топологий
Вот экспорт основных слоёв рекуррентной нейросети с одним RNN слоем.
using Flux
using ONNXNaiveNASflux
# Эта модель принимает 5 переменных, производит их обработку при помощи RNN и выдает 2 значения
rnn_model = Chain(
RNN(5 => 10), # Входы: 5 признаков, Внутренний слой: 10 признаков
Dense(10 => 2) # Выходной слой переводит состояние скрытого слоя в два класса/значения
)
# Определим размерность входного тензора:
# Переменных = 5, Размер батча = 1, Длина входной последовательности = 100
input_shape = (5, 1, 100)
ONNXNaiveNASflux.save("rnn_model.onnx", rnn_model, input_shape);
Запустим рекуррентную сеть из файла ONNX:
import ONNXRunTime as ORT
onnx_rnn_model = ORT.load_inference("rnn_model.onnx")
onnx_rnn_model
С этой информацией мы можем запустить рекуррентную нейросеть и прогнозировать выходные признаки для каждого отсчета входного временного ряда:
plot()
for x in 1:20
# Генерируем тестовую последовательность : 5 признаков на 100 отсчетов
raw_data = rand(Float32, 5, 100)
# Добавляем размерность с учетом батча (стандарт ONNX)
input_tensor = reshape(raw_data', 100, 1, 5)
outputs = onnx_rnn_model(Dict("data_0" => input_tensor))
onnx_output = outputs["dense_0"]
# Пересобираем данные в матрицу (100 строк, 2 колонки)
plot_data = reshape(onnx_output, 100, 2)
gr()
plot!(plot_data, title="Прогноз RNN для 20 случайных запусков",
label = (x==1 ? ["Выход 1" "Выход 2"] : false),
line = (x==1 ? (5,:solid) : (1,:dot)),
c = [:red :black],
xlabel="Временной шаг (отсчет)")
end
plot!()
Сверточные сети
Если соблюдать порядок аргументов, то несложно выполнить то же самое для свёрточных нейросетей:
using Flux
using ONNXNaiveNASflux
# Простая сверточная сеть
# Вход: 28x28 (1 канал), Выход: 10 классов (например, цифры 0-9)
cnn_model = Chain(
Conv((3, 3), 1 => 4, relu), # Свертка: фильтр 3x3, 1 входной канал, 4 выходных
Flux.flatten, # Выпрямление тензора в вектор
Dense(26 * 26 * 4 => 10) # Полносвязный слой для классификации
)
# Задаем размерность в формате Flux/Julia: (Width=28, Height=28, Channels=1, Batch=1)
input_shape_cnn = (28, 28, 1, 1)
# Экспортируем
ONNXNaiveNASflux.save("cnn_model.onnx", cnn_model, input_shape_cnn);
И запустим сохранённую в ONNX свёрточную нейросеть:
import ONNXRunTime as ORT
# Загружаем модель
onnx_cnn_model = ORT.load_inference("cnn_model.onnx")
# 1. Генерируем случайную "картинку" 28x28 с 1 каналом в формате Julia (W, H, C)
# В реальной задаче здесь будет реальное изображение
raw_image = rand(Float32, 28, 28, 1)
# 2. Переставляем оси для ONNX
# Из (W=28, H=28, C=1) делаем (C=1, H=28, W=28)
permuted_image = permutedims(raw_image, (3, 2, 1))
# Добавляем размерность батча на первую позицию -> получаем (Batch=1, C=1, H=28, W=28)
input_tensor = reshape(permuted_image, 1, 1, 28, 28)
# 3. Передаем в модель (имя входа для первого слоя обычно "data_0")
input_data = Dict("data_0" => input_tensor)
outputs = onnx_cnn_model(input_data)
# 4. Извлекаем выходные данные
# На выходе будет тензор с логитами (вероятностями) для 10 классов
onnx_output = outputs["dense_0"]
# 5. Принудительно приводим к одномерному вектору (10 элементов) для графика
# Это защитит нас от любой структуры батчей на выходе (будь то 1x10 или 10x1)
plot_data = reshape(onnx_output, :)
# Визуализация прогноза по 10 классам
gr()
bar(plot_data, title="Прогноз CNN для случайного изображения",
xticks=1:10, xlabel="Классы (0-9)", ylabel="Уверенность сети", legend=false)
Заключение
Мы разобрались с тем, как сохранить в ONNX нейросети трех основных видов. Не все конструкции из Flux экспортируются надежным образом, или вообще имеют поддержку в ONNXNaiveNASflux. Могут потребоваться эксперименты и упрощение вашей топологии.